feat(ml): optimized model inference (#31750)

* add immich_model

* graph per shape

* decouple ocr from ort

* prepared graphs

* use f16 for accelerators

* preload all shapes

* warm-up pass after preload

* keep precision for fp32 graphs

* ensure padding for text recognition

* optimize postprocessing for fp16 outputs

* handle raw ctc logits

* use f16 output for rknn

* mmap token embedding

* read device metadata

machine-learning/immich_ml/config.py

@@ -13,7 +13,7 @@ from rich.logging import RichHandler
 from uvicorn import Server
 from uvicorn.workers import UvicornWorker
 
-from .schemas import ModelOrganization, ModelPrecision
+from .schemas import ModelOrganization
 
 
 class ClipSettings(BaseModel):
@@ -29,6 +29,7 @@ class FacialRecognitionSettings(BaseModel):
 class OcrSettings(BaseModel):
     recognition: str | None = None
     detection: str | None = None
+    max_resolution: int = 736  # must match the server's OCR setting, or the preloaded model will not be used
 
 
 class PreloadModelData(BaseModel):
@@ -38,8 +39,8 @@ class PreloadModelData(BaseModel):
 
 
 class MaxBatchSize(BaseModel):
-    facial_recognition: int | None = None
-    ocr: int | None = None
+    facial_recognition: int = 4
+    ocr: int = 6
 
 
 def default_worker_timeout() -> int:
@@ -71,9 +72,7 @@ class Settings(BaseSettings):
     rknn: bool = True
     rknn_threads: int = 1
     preload: PreloadModelData | None = None
-    max_batch_size: MaxBatchSize | None = None
-    openvino_precision: ModelPrecision = ModelPrecision.FP32
-    rocm_precision: ModelPrecision = ModelPrecision.FP32
+    max_batch_size: MaxBatchSize = MaxBatchSize()
     model_organization: ModelOrganization = ModelOrganization.APP
     model_revision: str = "main"
 

machine-learning/immich_ml/main.py

@@ -91,6 +91,7 @@ async def preload_models(preload: PreloadModelData) -> None:
             model_name = model_name.strip()
             model = model_cache.get(model_name, model_type, model_task, **options)
             await load(model)
+            await attempt(model, model.build, MODEL_FILE_ERRORS)
 
     if preload.clip.textual is not None:
         await load_models(preload.clip.textual, ModelType.TEXTUAL, ModelTask.SEARCH)
@@ -117,6 +118,7 @@ async def preload_models(preload: PreloadModelData) -> None:
             preload.ocr.detection,
             ModelType.DETECTION,
             ModelTask.OCR,
+            maxResolution=preload.ocr.max_resolution,
         )
 
     if preload.ocr.recognition is not None:

machine-learning/immich_ml/models/base.py

@@ -15,6 +15,7 @@ from immich_ml.sessions.ort import OrtSession
 from ..config import clean_name, log, settings
 from ..schemas import ModelFormat, ModelIdentity, ModelSession, ModelTask, ModelType
 from ..sessions.ann import AnnSession
+from ..sessions.policy import ShapePolicy
 
 _IGNORED_PATTERNS: dict[ModelFormat, list[str]] = {
     ModelFormat.ONNX: ["*.armnn", "*.rknn"],
@@ -26,6 +27,10 @@ _IGNORED_PATTERNS: dict[ModelFormat, list[str]] = {
 class InferenceModel(ABC):
     depends: ClassVar[list[ModelIdentity]]
     identity: ClassVar[ModelIdentity]
+    # options a graph is built for, so one instance cannot serve two values of them
+    graph_options: ClassVar[tuple[str, ...]] = ()
+    # the shapes this model feeds, which pick the artifact and tell the engine what it may pin
+    shape_policy: ShapePolicy = ShapePolicy()
     threads: ClassVar[int] = 2  # how many threads one run takes for CPU inference
 
     def __init__(
@@ -78,6 +83,10 @@ class InferenceModel(ABC):
                 del self.session
                 self.loaded = False
 
+    def build(self) -> None:
+        self.load()
+        self.session.warm()
+
     def predict(self, *inputs: Any, **model_kwargs: Any) -> Any:
         self.load()
         return self._predict(*inputs, **model_kwargs)
@@ -115,23 +124,19 @@ class InferenceModel(ABC):
             case ModelFormat.ARMNN:
                 return AnnSession(self.model_path)
             case ModelFormat.ONNX:
-                return OrtSession(self.model_path, threads=self.threads)
+                return OrtSession(self.model_path, self.shape_policy, threads=self.threads)
             case ModelFormat.RKNN:
                 return rknn.RknnSession(self.model_path)
 
-    def model_path_for_format(self, model_format: ModelFormat) -> Path:
-        model_path_prefix = rknn.model_prefix if model_format == ModelFormat.RKNN else None
-        if model_path_prefix:
-            return self.model_dir / model_path_prefix / f"model.{model_format}"
-        return self.model_dir / f"model.{model_format}"
-
     @property
     def model_dir(self) -> Path:
         return self.cache_dir / self.model_type.value
 
     @property
     def model_path(self) -> Path:
-        return self.model_path_for_format(self.model_format)
+        if self.model_format == ModelFormat.RKNN:
+            return rknn.model_path(self.model_dir, self.shape_policy.label)
+        return self.model_dir / f"model.{self.model_format}"
 
     @property
     def model_task(self) -> ModelTask:
@@ -151,7 +156,8 @@ class InferenceModel(ABC):
 
     @property
     def _cache_dir_default(self) -> Path:
-        return settings.cache_folder / self.model_task.value / self.model_name
+        cache_dir = settings.cache_folder / self.model_task.value / self.model_name
+        return cache_dir if settings.legacy_models else cache_dir / settings.model_revision
 
     @property
     def cached(self) -> bool:

machine-learning/immich_ml/models/cache.py

@@ -17,7 +17,7 @@ class ModelCache:
         self, model_name: str, model_type: ModelType, model_task: ModelTask, ttl: int | None = None, **options: Any
     ) -> InferenceModel:
         model_cls = get_model_class(model_name, model_type, model_task)
-        key = (model_name, model_type, model_task)
+        key = (model_name, model_type, model_task, *(options.get(option) for option in model_cls.graph_options))
         if key not in self._models:
             self._models[key] = model_cls(model_name, **options)
         elif key not in self._expiries:

machine-learning/immich_ml/models/clip/textual.py

@@ -15,18 +15,33 @@ from immich_ml.models.transforms import clean_text, serialize_np_array
 from immich_ml.schemas import ModelSession, ModelTask, ModelType
 
 
+def _mapped(model_path: Path, name: str) -> NDArray[Any]:
+    import onnx
+
+    initializer = next(i for i in onnx.load(model_path, load_external_data=False).graph.initializer if i.name == name)
+    data = {entry.key: entry.value for entry in initializer.external_data}
+    dtype = onnx.helper.tensor_dtype_to_np_dtype(initializer.data_type)
+    offset, shape = int(data.get("offset", 0)), tuple(initializer.dims)
+    return np.memmap(model_path.parent / data["location"], dtype=dtype, mode="r", offset=offset, shape=shape)
+
+
 class BaseCLIPTextualEncoder(InferenceModel):
     depends = []
     identity = (ModelType.TEXTUAL, ModelTask.SEARCH)
     threads = 4  # lower search latency
 
     def _predict(self, inputs: str, language: str | None = None) -> str:
-        tokens = self.tokenize(inputs, language=language)
-        res: NDArray[np.float32] = self.session.run(None, tokens)[0][0]
+        tokens: dict[str, NDArray[Any]] = self.tokenize(inputs, language=language)
+        graph = self.session.for_shape(self.shape_policy.dims[0])
+        if self.embedding is not None:  # a graph that left its token table to the host
+            tokens["token_embeds"] = self.embedding[next(iter(tokens.values()))]
+        res: NDArray[np.float32] = graph.run(None, {node.name: tokens[node.name] for node in graph.get_inputs()})[0][0]
         return serialize_np_array(res)
 
     def _load(self) -> ModelSession:
         session = super()._load()
+        table = session.for_shape(self.shape_policy.dims[0]).get_metadata().get("embedding")
+        self.embedding = _mapped(self.model_dir / "model.onnx", table) if table is not None else None
         log.debug(f"Loading tokenizer for CLIP model '{self.model_name}'")
         self.tokenizer = self._load_tokenizer()
         tokenizer_kwargs: dict[str, Any] | None = self.text_cfg.get("tokenizer_kwargs")

machine-learning/immich_ml/models/clip/visual.py

@@ -19,7 +19,7 @@ from immich_ml.models.transforms import (
     serialize_np_array,
     to_numpy,
 )
-from immich_ml.schemas import ModelSession, ModelTask, ModelType
+from immich_ml.schemas import ModelGraph, ModelInput, ModelSession, ModelTask, ModelType
 
 
 class BaseCLIPVisualEncoder(InferenceModel):
@@ -28,11 +28,12 @@ class BaseCLIPVisualEncoder(InferenceModel):
 
     def _predict(self, inputs: Image.Image | bytes) -> str:
         image = decode_pil(inputs)
-        res: NDArray[np.float32] = self.session.run(None, self.transform(image))[0][0]
+        session = self.session.for_shape(self.shape_policy.dims[0])
+        res: NDArray[np.float32] = session.run(None, self.transform(session, image))[0][0]
         return serialize_np_array(res)
 
     @abstractmethod
-    def transform(self, image: Image.Image) -> dict[str, NDArray[np.float32]]:
+    def transform(self, session: ModelGraph, image: Image.Image) -> ModelInput:
         pass
 
     @property
@@ -69,11 +70,12 @@ class OpenClipVisualEncoder(BaseCLIPVisualEncoder):
 
         return super()._load()
 
-    def transform(self, image: Image.Image) -> dict[str, NDArray[np.float32]]:
+    def transform(self, session: ModelGraph, image: Image.Image) -> ModelInput:
         image = self._resize(image)
-        image_np = to_numpy(image)
-        image_np = normalize(image_np, self.mean, self.std)
-        return {"image": np.expand_dims(image_np.transpose(2, 0, 1), 0)}
+        if session.normalizes_input:
+            rgb: NDArray[np.uint8] = np.asarray(image if image.mode == "RGB" else image.convert("RGB"))
+            return {"image": rgb[None]}
+        return {"image": normalize(to_numpy(image), self.mean, self.std).transpose(2, 0, 1)[None]}
 
     def _resize(self, image: Image.Image) -> Image.Image:
         match self.preprocess_cfg.get("resize_mode", "shortest"):  # open_clip's default for older configs

machine-learning/immich_ml/models/facial_recognition/_ops.py

@@ -8,11 +8,11 @@ from functools import lru_cache
 
 import cv2
 import numpy as np
+from immich_model.constants import FACE_DETECTION_SIZE as DET_SIZE
 from numpy.typing import NDArray
 
 from immich_ml.models.transforms import ensure_dims
 
-DET_SIZE = 640
 ALIGNED_SIZE = 112
 
 # the FPN levels for which the fused detector emits a (scores, boxes, kps) triple,
@@ -69,6 +69,5 @@ def umeyama(src: NDArray[np.float32], dst: NDArray[np.float32]) -> NDArray[np.fl
     return np.hstack([scale * rotation, translation[:, None]], dtype=np.float32)
 
 
-def align_face(image: NDArray[np.uint8], kps: NDArray[np.float32]) -> NDArray[np.float32]:
-    matrix = umeyama(kps, ARCFACE_DST)
-    return cv2.warpAffine(image, matrix, (ALIGNED_SIZE, ALIGNED_SIZE)).astype(np.float32)
+def align_face(image: NDArray[np.uint8], kps: NDArray[np.float32], crop: NDArray[np.uint8]) -> None:
+    cv2.warpAffine(image, umeyama(kps, ARCFACE_DST), (ALIGNED_SIZE, ALIGNED_SIZE), dst=crop)

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

@@ -1,23 +1,30 @@
 import numpy as np
+from immich_model.constants import FACE_DETECTION_SIZE as DET_SIZE
 from numpy.typing import NDArray
 
 from immich_ml.models.base import InferenceModel
-from immich_ml.models.transforms import decode_pil, letterbox, normalize
-from immich_ml.schemas import FaceDetectionOutput, ModelTask, ModelType
+from immich_ml.models.transforms import decode_pil, letterbox, normalize, widen
+from immich_ml.schemas import FaceDetectionOutput, ModelTask, ModelType, Shape
+from immich_ml.sessions.policy import ShapePolicy
 
-from ._ops import DET_SIZE, decode_scrfd, nms
+from ._ops import decode_scrfd, nms
 
 
 class FaceDetector(InferenceModel):
     depends = []
     identity = (ModelType.DETECTION, ModelTask.FACIAL_RECOGNITION)
+    shape_policy = ShapePolicy(dims=(Shape(batch=1, height=DET_SIZE, width=DET_SIZE),))
 
     def _predict(self, inputs: NDArray[np.uint8] | bytes, minScore: float) -> FaceDetectionOutput:
+        session = self.session.for_shape(self.shape_policy.dims[0])
         canvas, scale = letterbox(decode_pil(inputs), DET_SIZE)
-        blob = normalize(canvas.astype(np.float32), mean=127.5, std=128).transpose(2, 0, 1)[None]
+        blob: NDArray[np.float32] | NDArray[np.uint8] = (
+            canvas[None]
+            if session.normalizes_input
+            else normalize(canvas.astype(np.float32), 127.5, 128).transpose(2, 0, 1)[None]
+        )
 
-        input_name = self.session.get_inputs()[0].name
-        heads = self.session.run(None, {input_name: blob})
+        heads = [widen(head) for head in session.run(None, {session.get_inputs()[0].name: blob})]
         scores, boxes, kps = decode_scrfd(heads, DET_SIZE)
 
         candidates = scores >= minScore

machine-learning/immich_ml/models/facial_recognition/recognition.py

@@ -1,26 +1,16 @@
-from pathlib import Path
 from typing import Any
 
 import numpy as np
-import onnx
-import onnxruntime as ort
 from numpy.typing import NDArray
-from onnx.tools.update_model_dims import update_inputs_outputs_dims
 from PIL import Image
 
-from immich_ml.config import log, settings
+from immich_ml.config import settings
 from immich_ml.models.base import InferenceModel
 from immich_ml.models.transforms import decode_pil, normalize, serialize_np_array
-from immich_ml.schemas import (
-    FaceDetectionOutput,
-    FacialRecognitionOutput,
-    ModelFormat,
-    ModelSession,
-    ModelTask,
-    ModelType,
-)
+from immich_ml.schemas import FaceDetectionOutput, FacialRecognitionOutput, ModelTask, ModelType, Shape
+from immich_ml.sessions.policy import ShapePolicy, batches, runs
 
-from ._ops import align_face
+from ._ops import ALIGNED_SIZE, align_face
 
 
 class FaceRecognizer(InferenceModel):
@@ -29,15 +19,8 @@ class FaceRecognizer(InferenceModel):
 
     def __init__(self, model_name: str, **model_kwargs: Any) -> None:
         super().__init__(model_name, **model_kwargs)
-        max_batch_size = settings.max_batch_size and settings.max_batch_size.facial_recognition
-        self.batch_size = max_batch_size if max_batch_size else self._batch_size_default
-
-    def _load(self) -> ModelSession:
-        session = self._make_session()
-        if (not self.batch_size or self.batch_size > 1) and str(session.get_inputs()[0].shape[0]) != "batch":
-            self._add_batch_axis(self.model_path)
-            session = self._make_session()
-        return session
+        sizes = batches(settings.max_batch_size.facial_recognition)
+        self.shape_policy = ShapePolicy(dims=tuple(Shape(batch) for batch in sizes))
 
     def _predict(
         self, inputs: NDArray[np.uint8] | bytes | Image.Image, faces: FaceDetectionOutput
@@ -45,20 +28,26 @@ class FaceRecognizer(InferenceModel):
         if faces["boxes"].shape[0] == 0:
             return []
         image = np.asarray(decode_pil(inputs), dtype=np.uint8)
-        crops = np.stack([align_face(image, kps).transpose(2, 0, 1) for kps in faces["landmarks"]])
-        embeddings = self._predict_batch(normalize(crops, mean=127.5, std=127.5))
+        landmarks = faces["landmarks"]
+        sizes = runs(len(landmarks), self.session.batches)
+        crops: NDArray[np.float32] | NDArray[np.uint8] = np.empty(
+            (len(landmarks), ALIGNED_SIZE, ALIGNED_SIZE, 3), dtype=np.uint8
+        )
+        for crop, kps in zip(crops, landmarks):
+            align_face(image, kps, crop)
+        if not self.session.for_shape(Shape(sizes[0])).normalizes_input:
+            crops = normalize(crops.transpose(0, 3, 1, 2).astype(np.float32, order="C"), mean=127.5, std=127.5)
+        embeddings = self._predict_batch(crops, sizes)
         return self.postprocess(faces, embeddings)
 
-    def _predict_batch(self, crops: NDArray[np.float32]) -> NDArray[np.float32]:
-        input_name = self.session.get_inputs()[0].name
-        if not self.batch_size or crops.shape[0] <= self.batch_size:
-            return self.session.run(None, {input_name: crops})[0]
-
-        batches = [
-            self.session.run(None, {input_name: crops[i : i + self.batch_size]})[0]
-            for i in range(0, crops.shape[0], self.batch_size)
-        ]
-        return np.concatenate(batches, axis=0)
+    def _predict_batch(self, crops: NDArray[np.float32] | NDArray[np.uint8], sizes: list[int]) -> NDArray[np.float32]:
+        embeddings = []
+        start = 0
+        for size in sizes:
+            session = self.session.for_shape(Shape(size))
+            embeddings.append(session.run(None, {session.get_inputs()[0].name: crops[start : start + size]})[0])
+            start += size
+        return embeddings[0] if len(embeddings) == 1 else np.concatenate(embeddings, axis=0)
 
     def postprocess(self, faces: FaceDetectionOutput, embeddings: NDArray[np.float32]) -> FacialRecognitionOutput:
         return [
@@ -69,24 +58,3 @@ class FaceRecognizer(InferenceModel):
             }
             for (x1, y1, x2, y2), embedding, score in zip(faces["boxes"], embeddings, faces["scores"])
         ]
-
-    def _add_batch_axis(self, model_path: Path) -> None:
-        log.debug(f"Adding batch axis to model {model_path}")
-        proto = onnx.load(model_path)
-        static_input_dims = [shape.dim_value for shape in proto.graph.input[0].type.tensor_type.shape.dim[1:]]
-        static_output_dims = [shape.dim_value for shape in proto.graph.output[0].type.tensor_type.shape.dim[1:]]
-        input_dims = {proto.graph.input[0].name: ["batch"] + static_input_dims}
-        output_dims = {proto.graph.output[0].name: ["batch"] + static_output_dims}
-        updated_proto = update_inputs_outputs_dims(proto, input_dims, output_dims)
-        onnx.save(updated_proto, model_path)
-
-    @property
-    def _batch_size_default(self) -> int | None:
-        providers = ort.get_available_providers()
-        if (
-            self.model_format == ModelFormat.ONNX
-            and "MIGraphXExecutionProvider" not in providers
-            and "OpenVINOExecutionProvider" not in providers
-        ):
-            return None
-        return 1

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

@@ -1,38 +1,83 @@
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from pathlib import Path
 from typing import Any, Self
 
 import numpy as np
 from numpy.typing import NDArray
 
+from immich_ml.models.transforms import widen
 
-def greedy(outputs: Sequence[NDArray[Any]]) -> tuple[NDArray[np.int32], NDArray[np.float32]]:
-    """Models with an in-graph argmax head emit (indices, probs); the rest emit raw logits."""
-    if len(outputs) == 2:
-        return outputs[0], outputs[1]
-    (probs,) = outputs
-    indices = probs.argmax(axis=2)
-    return indices.astype(np.int32), np.take_along_axis(probs, indices[:, :, None], axis=2)[:, :, 0]
+
+def picked(indices: NDArray[np.int32], confidence: NDArray[Any]) -> tuple[NDArray[np.int32], NDArray[np.float32]]:
+    return indices, widen(confidence)
+
+
+def kept(indices: NDArray[np.integer[Any]]) -> NDArray[np.bool_]:
+    """The steps a greedy CTC decode reads a character from: not blank, and not a repeat of the step before."""
+    keep = np.empty(indices.shape, dtype=bool)
+    keep[:, 0] = True
+    np.not_equal(indices[:, 1:], indices[:, :-1], out=keep[:, 1:])
+    keep &= indices != 0
+    return keep
+
+
+def logits(raw: NDArray[Any]) -> tuple[NDArray[np.int32], NDArray[np.float32]]:
+    steps = raw.reshape(-1, raw.shape[-1])  # a binary keeps a unit axis between the steps and the classes
+    indices = _argmax(steps)
+    keep = kept(indices.reshape(raw.shape[:2])).ravel()
+    rows = widen(steps[keep])  # a confidence is only read where a character is kept
+    best = rows[np.arange(len(rows)), indices[keep]]
+    # the winner's softmax, summing only the few classes within exp(-16) of it rather than every class
+    near = rows > (best - 16)[:, None]
+    counts = np.count_nonzero(near, axis=1)
+    confidence = np.zeros(len(steps), np.float32)
+    if len(rows):
+        weights = np.exp(rows[near] - np.repeat(best, counts))
+        confidence[keep] = 1 / np.add.reduceat(weights, np.cumsum(counts) - counts)
+    return indices.reshape(raw.shape[:2]).astype(np.int32), confidence.reshape(raw.shape[:2])
+
+
+def _argmax(steps: NDArray[Any]) -> NDArray[np.intp]:
+    if steps.dtype != np.float16:
+        return steps.argmax(axis=1)
+    bits = steps.view(np.int16)  # a half's bits order like its value where it is not negative
+    indices = bits.argmax(axis=1)
+    negative = np.flatnonzero(bits[np.arange(len(bits)), indices] < 0)  # a step with no logit above zero
+    if len(negative):
+        flipped = bits[negative] ^ ((bits[negative] >> 15) & np.int16(0x7FFF))  # negatives' order restored
+        indices[negative] = flipped.argmax(axis=1)
+    return indices
+
+
+def probabilities(probs: NDArray[Any]) -> tuple[NDArray[np.int32], NDArray[np.float32]]:
+    # probabilities are non-negative, so half-precision bits order like their values and need no widening to compare
+    indices = (probs.view(np.uint16) if probs.dtype == np.float16 else probs).argmax(axis=2)
+    return indices.astype(np.int32), widen(np.take_along_axis(probs, indices[:, :, None], axis=2)[:, :, 0])
+
+
+Greedy = Callable[..., tuple[NDArray[np.int32], NDArray[np.float32]]]
+GREEDY: dict[str, Greedy] = {"ctc_indices": picked, "ctc_logits": logits, "logits_softmax": probabilities}
 
 
 class CtcDecoder:
-    def __init__(self, charset: list[str]):
+    def __init__(self, charset: list[str], greedy: Greedy):
         self.charset = ["", *charset, " "]  # PP-OCR: blank at 0, space last
+        self.greedy = greedy
 
     @classmethod
-    def from_file(cls, charset_path: Path) -> Self:
+    def from_file(cls, charset_path: Path, greedy: Greedy) -> Self:
         if not charset_path.is_file():
             raise FileNotFoundError(f"Recognition charset not found: {charset_path}")
-        return cls(charset_path.read_text(encoding="utf-8").splitlines())
+        return cls(charset_path.read_text(encoding="utf-8").splitlines(), greedy)
 
     def __len__(self) -> int:
         return len(self.charset)
 
+    def __call__(self, outputs: Sequence[NDArray[Any]]) -> tuple[list[str], NDArray[np.float32]]:
+        return self.decode(*self.greedy(*outputs))
+
     def decode(self, indices: NDArray[np.int32], probs: NDArray[np.float32]) -> tuple[list[str], NDArray[np.float32]]:
-        keep = np.empty(indices.shape, dtype=bool)
-        keep[:, 0] = True
-        np.not_equal(indices[:, 1:], indices[:, :-1], out=keep[:, 1:])  # repeats before blanks
-        keep &= indices != 0
+        keep = kept(indices)
 
         scores: NDArray[np.float32] = np.where(keep, probs, 0).sum(1)
         scores /= np.maximum(keep.sum(1), 1)  # an all-blank row sums to 0, so it stays 0

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

@@ -1,29 +1,31 @@
-from functools import cached_property
+from collections.abc import Sequence
 from typing import Any
 
 import numpy as np
+from immich_model.constants import ocr_canvases
 from numpy.typing import NDArray
 from PIL import Image
-from rapidocr.inference_engine.base import FileInfo, InferSession
-from rapidocr.utils.download_file import DownloadFile, DownloadFileInput
-from rapidocr.utils.typings import EngineType, LangDet, OCRVersion, TaskType
-from rapidocr.utils.typings import ModelType as RapidModelType
 
-from immich_ml.config import log
-from immich_ml.models.base import InferenceModel
-from immich_ml.schemas import ModelFormat, ModelSession, ModelTask, ModelType
-from immich_ml.sessions.ort import OrtSession
+from immich_ml.models.transforms import widen
+from immich_ml.schemas import ModelGraph, ModelTask, ModelType, Shape
+from immich_ml.sessions.policy import ShapePolicy
 
+from .legacy import TextModel
 from .postprocess import DBPostProcess
 from .schemas import TextDetectionOutput
 
 
-class TextDetector(InferenceModel):
+class TextDetector(TextModel):
     depends = []
     identity = (ModelType.DETECTION, ModelTask.OCR)
+    graph_options = ("maxResolution",)
 
     def __init__(self, model_name: str, **model_kwargs: Any) -> None:
-        super().__init__(model_name, **model_kwargs, model_format=ModelFormat.ONNX)
+        super().__init__(model_name, **model_kwargs)
+        short_side = model_kwargs.get("maxResolution", 736)
+        # RKNPU ships a binary per short side, which the label picks
+        canvases = tuple(Shape(batch=1, **canvas) for canvas in ocr_canvases(short_side))
+        self.shape_policy = ShapePolicy(dims=canvases, label=f"res{short_side}")
         self.scale = np.float32(1.0 / 127.5)  # (x/255 - 0.5) / 0.5
         self._empty: TextDetectionOutput = {
             "boxes": np.empty(0, dtype=np.float32),
@@ -31,28 +33,6 @@ class TextDetector(InferenceModel):
         }
         self.postprocess = DBPostProcess(thresh=0.3, max_candidates=1000, unclip_ratio=1.6, use_dilation=True)
 
-    def _download(self) -> None:
-        model_info = InferSession.get_model_url(
-            FileInfo(
-                engine_type=EngineType.ONNXRUNTIME,
-                ocr_version=OCRVersion.PPOCRV5,
-                task_type=TaskType.DET,
-                lang_type=LangDet.CH,
-                model_type=RapidModelType.MOBILE if "mobile" in self.model_name else RapidModelType.SERVER,
-            )
-        )
-        download_params = DownloadFileInput(
-            file_url=model_info["model_dir"],
-            sha256=model_info["SHA256"],
-            save_path=self.model_path,
-            logger=log,
-        )
-        DownloadFile.run(download_params)
-
-    def _load(self) -> ModelSession:
-        # TODO: support other runtime sessions
-        return OrtSession(self.model_path)
-
     def _predict(
         self, inputs: Image.Image, maxResolution: int = 736, minScore: float = 0.5, scoreMode: str = "fast"
     ) -> TextDetectionOutput:
@@ -60,9 +40,8 @@ class TextDetector(InferenceModel):
         if width < 32 or height < 32:
             return self._empty
 
-        image, content = self._transform(inputs, maxResolution)
-        input_name = self.session.get_inputs()[0].name
-        probs = self.session.run(None, {input_name: image})[0][0]
+        session, image, content = self._transform(inputs, maxResolution)
+        probs = widen(session.run(None, {session.get_inputs()[0].name: image})[0][0])
         if probs.ndim == 3:
             probs = probs[0]
         boxes, scores = self.postprocess(probs[: content[0], : content[1]], (height, width), minScore, scoreMode)
@@ -71,27 +50,46 @@ class TextDetector(InferenceModel):
         order = self.reading_order(boxes)
         return {"boxes": boxes[order], "scores": scores[order]}
 
-    @cached_property
-    def raw_input(self) -> bool:
-        return self.session.get_inputs()[0].shape[-1] == 3  # NHWC models handle normalization and transpose internally
-
-    def _transform(self, img: Image.Image, max_resolution: int) -> tuple[NDArray[Any], tuple[int, int]]:
+    def _transform(self, img: Image.Image, max_resolution: int) -> tuple[ModelGraph, NDArray[Any], tuple[int, int]]:
         ratio = min(max_resolution / min(img.height, img.width), 1.0)
-        resize_h = max(self._round32(img.height * ratio), 32)
-        resize_w = max(self._round32(img.width * ratio), 32)
-        resized = img.resize((resize_w, resize_h), resample=Image.Resampling.LANCZOS)
+        content_h = max(self._round32(img.height * ratio), 32)
+        content_w = max(self._round32(img.width * ratio), 32)
+        canvas = self._canvas(self.session.shapes, content_h, content_w)
+        session = self.session.for_shape(Shape(batch=1, height=canvas[0], width=canvas[1]))
+        # a compiled canvas may be smaller than the image was sized for, so it shrinks again
+        scale = min(canvas[0] / content_h, canvas[1] / content_w, 1.0)
+        content_h = min(max(self._round32(content_h * scale), 32), canvas[0])
+        content_w = min(max(self._round32(content_w * scale), 32), canvas[1])
+        resized = img.resize((content_w, content_h), resample=Image.Resampling.LANCZOS)
         if resized.mode != "RGB":
             resized = resized.convert("RGB")
         array = np.asarray(resized)
-        if self.raw_input:
-            return array[None], (resize_h, resize_w)
+        content = (content_h, content_w)
+        height, width = canvas
+        if session.normalizes_input:
+            if canvas == content:
+                return session, array[None], content
+            image = np.zeros((1, height, width, 3), dtype=np.uint8)  # what is left of the canvas stays black
+            image[0, :content_h, :content_w] = array
+            return session, image, content
         # reverse plane order gets the BGR swap and the CHW transpose for free
-        out = np.empty((1, 3, resize_h, resize_w), dtype=np.float32)
+        out = np.empty((1, 3, height, width), dtype=np.float32)
+        if canvas != content:  # -1.0 is what the graph would have made of a black pixel
+            out[0, :, content_h:] = -1.0
+            out[0, :, :content_h, content_w:] = -1.0
         for channel in range(3):
-            plane = out[0, 2 - channel]
+            plane = out[0, 2 - channel, :content_h, :content_w]
             np.multiply(array[:, :, channel], self.scale, out=plane)
             plane -= 1.0
-        return out, (resize_h, resize_w)
+        return session, out, content
+
+    @staticmethod
+    def _canvas(shapes: Sequence[Shape], height: int, width: int) -> tuple[int, int]:
+        return min(
+            ((shape.height, shape.width) for shape in shapes if shape.height and shape.width),
+            key=lambda canvas: (-min(canvas[0] / height, canvas[1] / width, 1.0), canvas[0] * canvas[1]),
+            default=(height, width),
+        )
 
     @staticmethod
     def _round32(value: float) -> int:

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

@@ -0,0 +1,42 @@
+from typing import Any
+
+from immich_ml.config import log, settings
+from immich_ml.models.base import InferenceModel
+from immich_ml.schemas import ModelFormat, ModelType
+
+
+# TODO: Remove once everything uses the new model graphs
+class TextModel(InferenceModel):
+    def __init__(self, model_name: str, **model_kwargs: Any) -> None:
+        if settings.legacy_models:
+            model_kwargs["model_format"] = ModelFormat.ONNX  # the older exports come in no other format
+        super().__init__(model_name, **model_kwargs)
+
+    def download(self) -> None:
+        if not settings.legacy_models or self.cached:
+            return super().download()
+        # here, as only an install still on the older exports has a use for it
+        from rapidocr.inference_engine.base import FileInfo, InferSession
+        from rapidocr.utils.download_file import DownloadFile, DownloadFileInput
+        from rapidocr.utils.typings import EngineType, LangDet, LangRec, OCRVersion, TaskType
+        from rapidocr.utils.typings import ModelType as RapidModelType
+
+        detects = self.model_type == ModelType.DETECTION
+        language = self.model_name.split("__")[0] if "__" in self.model_name else "CH"
+        model_info = InferSession.get_model_url(
+            FileInfo(
+                engine_type=EngineType.ONNXRUNTIME,
+                ocr_version=OCRVersion.PPOCRV5,
+                task_type=TaskType.DET if detects else TaskType.REC,
+                lang_type=LangDet.CH if detects else LangRec[language],
+                model_type=RapidModelType.MOBILE if "mobile" in self.model_name else RapidModelType.SERVER,
+            )
+        )
+        log.info(
+            f"Downloading {self.model_type} model '{self.model_name}' to {self.model_path}. This may take a while."
+        )
+        DownloadFile.run(
+            DownloadFileInput(
+                file_url=model_info["model_dir"], sha256=model_info["SHA256"], save_path=self.model_path, logger=log
+            )
+        )

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

@@ -1,77 +1,57 @@
 import math
-from functools import cached_property
+from bisect import bisect_left
 from typing import Any
 
 import cv2
 import numpy as np
+from immich_model.constants import OCR_RECOGNITION_WIDTHS
 from numpy.typing import NDArray
 from PIL import Image
-from rapidocr import LangRec
-from rapidocr.inference_engine.base import FileInfo, InferSession
-from rapidocr.utils.download_file import DownloadFile, DownloadFileInput
-from rapidocr.utils.typings import EngineType, OCRVersion, TaskType
-from rapidocr.utils.typings import ModelType as RapidModelType
 
-from immich_ml.config import log, settings
-from immich_ml.models.base import InferenceModel
-from immich_ml.schemas import ModelFormat, ModelSession, ModelTask, ModelType
-from immich_ml.sessions.ort import OrtSession
+from immich_ml.config import settings
+from immich_ml.schemas import ModelGraph, ModelSession, ModelTask, ModelType, Shape
+from immich_ml.sessions.policy import ShapePolicy, batches, runs
 
-from .ctc import CtcDecoder, greedy
+from .ctc import GREEDY, CtcDecoder, probabilities
+from .legacy import TextModel
 from .schemas import TextDetectionOutput, TextRecognitionOutput
 
 REC_HEIGHT = 48
-REC_BASE_RATIO = 320 / REC_HEIGHT  # PP-OCR's rec_image_shape floor
 SCALE = np.float32(1.0 / 127.5)
 
 
-class TextRecognizer(InferenceModel):
+class TextRecognizer(TextModel):
     depends = [(ModelType.DETECTION, ModelTask.OCR)]
     identity = (ModelType.RECOGNITION, ModelTask.OCR)
 
     def __init__(self, model_name: str, **model_kwargs: Any) -> None:
-        self.language = LangRec[model_name.split("__")[0]] if "__" in model_name else LangRec.CH
         self._empty: TextRecognitionOutput = {
             "box": np.empty(0, dtype=np.float32),
             "boxScore": np.empty(0, dtype=np.float32),
             "text": [],
             "textScore": np.empty(0, dtype=np.float32),
         }
-        super().__init__(model_name, **model_kwargs, model_format=ModelFormat.ONNX)
-        max_batch_size = settings.max_batch_size and settings.max_batch_size.ocr
-        self.batch_size = max_batch_size if max_batch_size else 6
-
-    def _download(self) -> None:
-        model_info = InferSession.get_model_url(
-            FileInfo(
-                engine_type=EngineType.ONNXRUNTIME,
-                ocr_version=OCRVersion.PPOCRV5,
-                task_type=TaskType.REC,
-                lang_type=self.language,
-                model_type=RapidModelType.MOBILE if "mobile" in self.model_name else RapidModelType.SERVER,
-            )
+        super().__init__(model_name, **model_kwargs)
+        self.decoder: CtcDecoder | None = None
+        sizes = batches(settings.max_batch_size.ocr)
+        self.shape_policy = ShapePolicy(
+            dims=tuple(Shape(batch, width=width) for width in OCR_RECOGNITION_WIDTHS for batch in sizes)
         )
-        download_params = DownloadFileInput(
-            file_url=model_info["model_dir"],
-            sha256=model_info["SHA256"],
-            save_path=self.model_path,
-            logger=log,
-        )
-        DownloadFile.run(download_params)
 
     def _load(self) -> ModelSession:
-        # TODO: support other runtimes
-        session = OrtSession(self.model_path)
-        self.decoder = (
-            CtcDecoder(character.splitlines())
-            if (character := session.get_metadata().get("character")) is not None
-            else CtcDecoder.from_file(self.model_dir / "charset.txt")
-        )
+        session = super()._load()
+        # the widths it was compiled for, narrowest first: none where it takes any
+        self.widths = tuple(sorted({shape.width for shape in session.shapes if shape.width is not None}))
         return session
 
-    @cached_property
-    def raw_input(self) -> bool:
-        return self.session.get_inputs()[0].shape[-1] == 3  # NHWC models handle normalization and transpose internally
+    def _decoder(self, session: ModelGraph) -> CtcDecoder:
+        character = session.get_metadata().get("character")
+        greedy = GREEDY.get(session.get_outputs()[0].name, probabilities)
+        return (
+            CtcDecoder(character.splitlines(), greedy)
+            if character is not None
+            else CtcDecoder.from_file(self.model_dir / "charset.txt", greedy)
+        )
 
     def _predict(self, img: Image.Image, texts: TextDetectionOutput, minScore: float = 0.9) -> TextRecognitionOutput:
         boxes, box_scores = texts["boxes"], texts["scores"]
@@ -85,28 +65,35 @@ class TextRecognizer(InferenceModel):
 
         text_list: list[str] = [""] * len(order)
         score_list = np.zeros(len(order), dtype=np.float32)
-        input_name = self.session.get_inputs()[0].name
-        for start in range(0, len(order), self.batch_size):
-            chunk = order[start : start + self.batch_size]
-            width = int(REC_HEIGHT * max(REC_BASE_RATIO, ratios[chunk[-1]]))
-            images: NDArray[Any] = (
-                np.full((len(chunk), REC_HEIGHT, width, 3), 127, dtype=np.uint8)
-                if self.raw_input
-                else np.zeros((len(chunk), 3, REC_HEIGHT, width), dtype=np.float32)
-            )
+        sized = runs(len(order), self.session.batches)
+        bounds = [sum(sized[:index]) for index in range(len(sized))]
+        chunked = [order[start : start + size] for start, size in zip(bounds, sized)]
+        batch_widths = [self._width(int(REC_HEIGHT * ratios[chunk[-1]])) for chunk in chunked]
+        graphs = [self.session.for_shape(Shape(len(chunk), width=width)) for chunk, width in zip(chunked, batch_widths)]
+        raw = graphs[0].normalizes_input
+        # one buffer per request, sized for the largest batch and reshaped for each
+        sizes = [len(chunk) * REC_HEIGHT * width * 3 for chunk, width in zip(chunked, batch_widths)]
+        buffer = np.empty(max(sizes), dtype=np.uint8 if raw else np.float32)
+        for chunk, width, size, session in zip(chunked, batch_widths, sizes, graphs):
+            shape = (len(chunk), REC_HEIGHT, width, 3) if raw else (len(chunk), 3, REC_HEIGHT, width)
+            images: NDArray[Any] = buffer[:size].reshape(shape)
             for i, index in enumerate(chunk):
                 crop = self._crop(img, coeffs[index], widths[index], heights[index], upright[index])
                 resized_w = max(1, min(width, math.ceil(REC_HEIGHT * ratios[index])))
-                if self.raw_input:
+                if raw:
                     cv2.resize(crop, (resized_w, REC_HEIGHT), dst=images[i, :, :resized_w])
+                    images[i, :, resized_w:] = 127  # the pad the crop is read against, over whatever was there
                     continue
                 resized = cv2.resize(crop, (resized_w, REC_HEIGHT))
                 view = images[i, :, :, :resized_w]
                 np.multiply(resized.transpose(2, 0, 1)[::-1], SCALE, out=view)  # [::-1] is the RGB -> BGR swap
                 view -= 1.0
+                images[i, :, :, resized_w:] = 0
 
-            out_indices, out_probs = greedy(self.session.run(None, {input_name: images}))
-            chunk_texts, chunk_scores = self.decoder.decode(out_indices, out_probs)
+            outputs = session.run(None, {session.get_inputs()[0].name: images})
+            if self.decoder is None:
+                self.decoder = self._decoder(session)
+            chunk_texts, chunk_scores = self.decoder(outputs)
             for index, text, score in zip(chunk, chunk_texts, chunk_scores):
                 text_list[index] = text
                 score_list[index] = score
@@ -123,6 +110,11 @@ class TextRecognizer(InferenceModel):
             "textScore": score_list[valid],
         }
 
+    def _width(self, width: int) -> int:
+        if not self.widths:  # room past the text in proportion to it, as the recognizer reads a line worse without
+            return min(OCR_RECOGNITION_WIDTHS[-1], max(OCR_RECOGNITION_WIDTHS[0], math.ceil(width * 1.25)))
+        return self.widths[min(bisect_left(self.widths, width), len(self.widths) - 1)]
+
     def _crop_geometry(
         self, boxes: NDArray[np.float32]
     ) -> tuple[NDArray[np.int32], NDArray[np.int32], NDArray[np.float32]]:

machine-learning/immich_ml/models/transforms.py

@@ -1,5 +1,6 @@
 import string
 from io import BytesIO
+from typing import Any
 
 import cv2
 import numpy as np
@@ -43,10 +44,6 @@ def get_pil_resampling(resample: str) -> Image.Resampling:
     return _PIL_RESAMPLING_METHODS[resample.lower()]
 
 
-def pil_to_cv2(image: Image.Image) -> NDArray[np.uint8]:
-    return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)  # type: ignore
-
-
 def decode_pil(image_bytes: bytes | Image.Image | NDArray[np.uint8]) -> Image.Image:
     image: Image.Image
     match image_bytes:
@@ -75,6 +72,13 @@ def serialize_np_array(arr: NDArray[np.float32]) -> str:
     return orjson.dumps(arr, option=orjson.OPT_SERIALIZE_NUMPY).decode()
 
 
+def widen(array: NDArray[Any]) -> NDArray[Any]:
+    """Single precision from half through OpenCV, as numpy converts one value at a time without native half support."""
+    if array.dtype != np.float16:
+        return array
+    return cv2.convertFp16(array.reshape(-1, array.shape[-1]).view(np.int16)).reshape(array.shape)  # int16 carries half
+
+
 def letterbox(image: NDArray[np.uint8] | Image.Image, size: int) -> tuple[NDArray[np.uint8], float]:
     if isinstance(image, Image.Image):
         image = np.asarray(image)

machine-learning/immich_ml/schemas.py

@@ -1,6 +1,7 @@
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
+from dataclasses import asdict, dataclass
 from enum import Enum
-from typing import Any, Literal, Protocol, TypeVar
+from typing import Any, Literal, Protocol, TypeAlias, TypeVar
 
 import numpy as np
 import numpy.typing as npt
@@ -44,11 +45,6 @@ class ModelSource(StrEnum):
     PADDLE = "paddle"
 
 
-class ModelPrecision(StrEnum):
-    FP16 = "FP16"
-    FP32 = "FP32"
-
-
 class ModelOrganization(StrEnum):
     APP = "immich-app"
     TESTING = "immich-testing"
@@ -62,16 +58,29 @@ class SessionNode(Protocol):
     def name(self) -> str: ...
 
     @property
-    def shape(self) -> tuple[int, ...]: ...
+    def shape(self) -> tuple[int | str, ...]: ...  # ORT names a symbolic dim rather than sizing it
 
 
-class ModelSession(Protocol):
+ModelTensor: TypeAlias = npt.NDArray[np.float32] | npt.NDArray[np.int32] | npt.NDArray[np.uint8]
+ModelInput = Mapping[str, ModelTensor]
+
+
+@dataclass(frozen=True)
+class Shape:
+    batch: int
+    height: int | None = None
+    width: int | None = None
+
+    @property
+    def pins(self) -> dict[str, int]:
+        return {name: size for name, size in asdict(self).items() if size is not None}
+
+
+class ModelGraph(Protocol):
     def run(
         self,
         output_names: list[str] | None,
-        input_feed: dict[str, npt.NDArray[np.float32]]
-        | dict[str, npt.NDArray[np.int32]]
-        | dict[str, npt.NDArray[np.uint8]],
+        input_feed: ModelInput,
         run_options: Any = None,
     ) -> list[npt.NDArray[np.float32]]: ...
 
@@ -81,6 +90,21 @@ class ModelSession(Protocol):
 
     def get_metadata(self) -> dict[str, str]: ...
 
+    @property
+    def normalizes_input(self) -> bool: ...
+
+
+class ModelSession(Protocol):
+    @property
+    def shapes(self) -> tuple[Shape, ...]: ...
+
+    @property
+    def batches(self) -> tuple[int, ...]: ...
+
+    def for_shape(self, shape: Shape) -> ModelGraph: ...
+
+    def warm(self) -> None: ...
+
 
 class FaceDetectionOutput(TypedDict):
     boxes: npt.NDArray[np.float32]

machine-learning/immich_ml/sessions/ann/__init__.py

@@ -7,7 +7,7 @@ import numpy as np
 from numpy.typing import NDArray
 
 from immich_ml.config import log, settings
-from immich_ml.schemas import SessionNode
+from immich_ml.schemas import ModelInput, SessionNode, Shape
 
 from .loader import Ann
 
@@ -29,6 +29,9 @@ class AnnSession:
             fp16=settings.ann_fp16_turbo,
         )
         log.info("Loaded ANN model with ID %d", self.model)
+        batch = self.ann.input_shapes[self.model][0][0]
+        self.shapes = (Shape(batch),)  # the artifact is compiled for one shape
+        self.batches = (batch,)
 
     def __del__(self) -> None:
         self.ann.unload(self.model)
@@ -46,10 +49,20 @@ class AnnSession:
     def get_metadata(self) -> dict[str, str]:
         return {}
 
+    @property
+    def normalizes_input(self) -> bool:
+        return False  # no ARM-NN artifact is built from a graph that carries its own preprocessing
+
+    def for_shape(self, shape: Shape) -> AnnSession:
+        return self
+
+    def warm(self) -> None:
+        pass  # the artifact is loaded whole when the session opens
+
     def run(
         self,
         output_names: list[str] | None,
-        input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]] | dict[str, NDArray[np.uint8]],
+        input_feed: ModelInput,
         run_options: Any = None,
     ) -> list[NDArray[np.float32]]:
         inputs: list[NDArray[np.float32]] = [np.ascontiguousarray(v) for v in input_feed.values()]

machine-learning/immich_ml/sessions/ort.py

@@ -1,77 +1,56 @@
 from __future__ import annotations
 
 import ctypes
+import pickle
 import platform
+import subprocess
 import sys
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from functools import cache, cached_property
 from pathlib import Path
 from threading import Lock
-from typing import Any
+from typing import Any, NamedTuple
 
 import numpy as np
 import onnxruntime as ort
+from immich_model.runtime import RewriteContext, RewritePlan, plan_rewrites
 from numpy.typing import NDArray
+from onnxruntime.capi.onnxruntime_pybind11_state import InvalidProtobuf
+from pydantic import BaseModel
 
-from immich_ml.models.constants import SUPPORTED_PROVIDERS
-from immich_ml.schemas import ModelPrecision, SessionNode
+from immich_ml.schemas import ModelInput, SessionNode, Shape
 
 from ..config import log, settings
+from .policy import ShapePolicy
 
-MigraphxInputSignature = tuple[tuple[str, str, tuple[int, ...]], ...]
-
-_migraphx_registry_lock = Lock()
-_migraphx_model_locks: dict[str, Lock] = {}
-_migraphx_compiled_inputs: set[tuple[str, MigraphxInputSignature]] = set()
+# the one provider that handles a free dim well; the rest miscompile it, recompile inside the run,
+# or partition around it, so they are handed one graph per shape instead
+DYNAMIC_PROVIDERS = frozenset({"CPUExecutionProvider"})
 
 
-def _migraphx_get_model_lock(model_key: str) -> Lock:
-    with _migraphx_registry_lock:
-        lock = _migraphx_model_locks.get(model_key)
-        if lock is None:
-            lock = Lock()
-            _migraphx_model_locks[model_key] = lock
-        return lock
+def _label(pins: Mapping[str, int]) -> str:
+    return "_".join(f"{name}{size}" for name, size in pins.items())
 
 
-def _migraphx_has_compiled_input(key: tuple[str, MigraphxInputSignature]) -> bool:
-    with _migraphx_registry_lock:
-        return key in _migraphx_compiled_inputs
+UNREADABLE = 3  # how the preparing child says its source does not parse
+DTYPES = {
+    "tensor(float)": np.float32,
+    "tensor(float16)": np.float16,
+    "tensor(uint8)": np.uint8,
+    "tensor(int32)": np.int32,
+    "tensor(int64)": np.int64,
+}
 
 
-def _migraphx_mark_compiled_input(key: tuple[str, MigraphxInputSignature]) -> None:
-    with _migraphx_registry_lock:
-        _migraphx_compiled_inputs.add(key)
-
-
-def _migraphx_input_signature(input_feed: dict[str, NDArray[Any]]) -> MigraphxInputSignature:
-    return tuple((name, str(value.dtype), tuple(value.shape)) for name, value in sorted(input_feed.items()))
-
-
-class OrtSession:
-    session: ort.InferenceSession
-
-    def __init__(
-        self,
-        model_path: Path | str,
-        providers: list[str] | None = None,
-        provider_options: list[dict[str, Any]] | None = None,
-        sess_options: ort.SessionOptions | None = None,
-        threads: int = 2,
-    ):
-        self.model_path = Path(model_path)
-        self.threads = threads
-        self.providers = providers if providers is not None else self._providers_default
-        self.disabled_optimizers = _disabled_optimizers_default(self.providers)
-        log.debug(f"Setting disabled_optimizers to {self.disabled_optimizers}")
-        self.provider_options = provider_options if provider_options is not None else self._provider_options_default
-        self.sess_options = sess_options if sess_options is not None else self._sess_options_default
-        self.session = ort.InferenceSession(
-            self.model_path.as_posix(),
-            providers=self.providers,
-            provider_options=self.provider_options,
-            sess_options=self.sess_options,
-            disabled_optimizers=self.disabled_optimizers,
-        )
+class OrtGraph:
+    def __init__(self, spec: GraphSpec) -> None:
+        self.session = spec.session(prepared(spec))
+        image = self.session.get_inputs()[0]
+        # whether the graph scales and shifts the image itself, so the host hands over raw pixels
+        self.normalizes_input = image.type == "tensor(uint8)"
+        # CoreML's rewrite moves the layout change out of the graph, so whoever applied it performs it
+        self.channels_first = self.normalizes_input and len(image.shape) == 4 and image.shape[1] == 3
 
     def get_inputs(self) -> Sequence[SessionNode]:
         inputs: Sequence[SessionNode] = self.session.get_inputs()
@@ -88,52 +67,91 @@ class OrtSession:
     def run(
         self,
         output_names: list[str] | None,
-        input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]] | dict[str, NDArray[np.uint8]],
+        input_feed: ModelInput,
         run_options: Any = None,
     ) -> list[NDArray[np.float32]]:
-        if "MIGraphXExecutionProvider" in self.providers:
-            model_key = self.model_path.resolve().as_posix()
-            input_key = (model_key, _migraphx_input_signature(input_feed))
-            if not _migraphx_has_compiled_input(input_key):
-                model_lock = _migraphx_get_model_lock(model_key)
-                with model_lock:
-                    if not _migraphx_has_compiled_input(input_key):
-                        outputs: list[NDArray[np.float32]] = self.session.run(output_names, input_feed, run_options)
-                        _migraphx_mark_compiled_input(input_key)
-                        return outputs
-
-            outputs = self.session.run(output_names, input_feed, run_options)
-            return outputs
-
-        outputs = self.session.run(output_names, input_feed, run_options)
+        if self.channels_first:
+            input_feed = {name: frames.transpose(0, 3, 1, 2) for name, frames in input_feed.items()}
+        outputs: list[NDArray[Any]] = self.session.run(output_names, input_feed, run_options)
         return outputs
 
-    @property
-    def providers(self) -> list[str]:
-        return self._providers
 
-    @providers.setter
-    def providers(self, providers: list[str]) -> None:
-        log.info(f"Setting execution providers to {providers}, in descending order of preference")
-        self._providers = providers
+class Device(NamedTuple):
+    kind: str  # what the provider compiles for, which names the directory
+    version: str  # what the facts go stale under
+
+
+@dataclass(frozen=True)
+class GraphSpec:
+    model_path: Path
+    pins: Mapping[str, int]
+    overrides: Sequence[tuple[str, int]]
+    providers: list[str]
+    disabled_optimizers: list[str]
+    threads: int = 2  # for the CPU, where they are not configured
 
     @property
-    def _providers_default(self) -> list[str]:
-        available_providers = set(ort.get_available_providers())
-        log.debug(f"Available ORT providers: {available_providers}")
-        return [provider for provider in SUPPORTED_PROVIDERS if provider in available_providers]
+    def provider(self) -> str:
+        return self.providers[0]  # the rest take only what it cannot run
+
+    @cached_property
+    def directory(self) -> Path:
+        provider = self.provider.removesuffix("ExecutionProvider").lower()
+        kind = self.device.kind if self.device else ""
+        return self.model_path.parent / provider / kind / _label(self.pins)
+
+    @cached_property
+    def device(self) -> Device | None:
+        match self.provider:
+            case "CPUExecutionProvider":
+                return Device(platform.machine(), _cpu())  # a prepack built for other instruction sets fails at runtime
+            case "OpenVINOExecutionProvider" if self.openvino_device.startswith("GPU"):
+                return _intel_gpu(self.openvino_device)
+            case "MIGraphXExecutionProvider":
+                return _amd_gpu(int(settings.device_id))
+        return None
+
+    @cached_property
+    def openvino_device(self) -> str:
+        device_ids: list[str] = ort.capi._pybind_state.get_available_openvino_device_ids()
+        return f"GPU.{settings.device_id}" if any(d.startswith("GPU") for d in device_ids) else "CPU"
+
+    @cached_property
+    def manifest(self) -> Path:
+        return self.directory / "manifest.json"
 
     @property
+    def plan(self) -> RewritePlan:
+        return _plan(self.provider)
+
+    @cached_property
+    def half(self) -> bool:
+        return not (settings.legacy_models or self.provider == "CPUExecutionProvider")
+
+    @cached_property
+    def facts(self) -> Facts:
+        return Facts(
+            half=self.half,
+            onnxruntime=ort.__version__,
+            # workers on devices of one kind share what is prepared
+            options={key: value for key, value in self.provider_options[0].items() if not key.startswith("device")},
+            overrides=list(self.overrides),
+            disabled_optimizers=self.disabled_optimizers,
+            rewrite=None if settings.legacy_models else self.plan.digest,  # which were never tested with the rewriter
+            device=self.device.version if self.device else None,
+        )
+
+    def session(self, graph: Path, sess_options: ort.SessionOptions | None = None) -> ort.InferenceSession:
+        return ort.InferenceSession(
+            graph.as_posix(),
+            providers=self.providers,
+            provider_options=self.provider_options,
+            sess_options=sess_options or self.sess_options(),
+            disabled_optimizers=self.disabled_optimizers,
+        )
+
+    @cached_property
     def provider_options(self) -> list[dict[str, Any]]:
-        return self._provider_options
-
-    @provider_options.setter
-    def provider_options(self, provider_options: list[dict[str, Any]]) -> None:
-        log.debug(f"Setting execution provider options to {provider_options}")
-        self._provider_options = provider_options
-
-    @property
-    def _provider_options_default(self) -> list[dict[str, Any]]:
         provider_options = []
         for provider in self.providers:
             match provider:
@@ -142,80 +160,245 @@ class OrtSession:
                 case "CUDAExecutionProvider":
                     options = {"arena_extend_strategy": "kSameAsRequested", "device_id": settings.device_id}
                 case "MIGraphXExecutionProvider":
-                    migraphx_dir = self.model_path.parent / "migraphx"
-                    # MIGraphX does not create the underlying folder and will crash if it does not exist
-                    migraphx_dir.mkdir(parents=True, exist_ok=True)
-                    options = {
-                        "device_id": settings.device_id,
-                        "migraphx_model_cache_dir": migraphx_dir.as_posix(),
-                        "migraphx_fp16_enable": "1" if settings.rocm_precision == ModelPrecision.FP16 else "0",
-                    }
+                    options = {"device_id": settings.device_id, "migraphx_model_cache_dir": self.directory.as_posix()}
                 case "OpenVINOExecutionProvider":
-                    device_ids: list[str] = ort.capi._pybind_state.get_available_openvino_device_ids()
-                    # Check for available devices, preferring GPU over CPU
-                    gpu_devices = [d for d in device_ids if d.startswith("GPU")]
-                    if gpu_devices:
-                        device_type = f"GPU.{settings.device_id}"
-                        log.debug(f"OpenVINO: Using GPU device {device_type}")
-                    else:
-                        device_type = "CPU"
-                        log.debug("OpenVINO: No GPU found, using CPU")
-                    options = {
-                        "device_type": device_type,
-                        "precision": settings.openvino_precision.value,
-                        "cache_dir": (self.model_path.parent / "openvino").as_posix(),
-                    }
+                    options = {"device_type": self.openvino_device, "cache_dir": self.directory.as_posix()}
+                    if not self.half:  # the GPU would otherwise run an fp32 graph at fp16
+                        options["precision"] = "FP32"
                 case "CoreMLExecutionProvider":
                     options = {
                         "ModelFormat": "MLProgram",
                         "MLComputeUnits": "ALL",
                         "SpecializationStrategy": "FastPrediction",
                         "AllowLowPrecisionAccumulationOnGPU": "1",
-                        "ModelCacheDirectory": (self.model_path.parent / "coreml").as_posix(),
+                        "ModelCacheDirectory": self.directory.as_posix(),
                     }
                 case _:
                     options = {}
             provider_options.append(options)
+        log.debug(f"Setting execution provider options to {provider_options}")
         return provider_options
 
-    @property
     def sess_options(self) -> ort.SessionOptions:
-        return self._sess_options
-
-    @sess_options.setter
-    def sess_options(self, sess_options: ort.SessionOptions) -> None:
-        log.debug(f"Setting execution_mode to {sess_options.execution_mode.name}")
-        log.debug(f"Setting inter_op_num_threads to {sess_options.inter_op_num_threads}")
-        log.debug(f"Setting intra_op_num_threads to {sess_options.intra_op_num_threads}")
-        self._sess_options = sess_options
-
-    @property
-    def _sess_options_default(self) -> ort.SessionOptions:
         sess_options = ort.SessionOptions()
+        sess_options.enable_cpu_mem_arena = settings.model_arena
         # some CPUs slow down many times over on subnormal operands, and ORT clears the flush in its threads otherwise
         sess_options.add_session_config_entry("session.set_denormal_as_zero", "1")
-        sess_options.enable_cpu_mem_arena = settings.model_arena
 
         # avoid thread contention between models
         # Set inter_op threads
         if settings.model_inter_op_threads > 0:
             sess_options.inter_op_num_threads = settings.model_inter_op_threads
         # these defaults work well for CPU, but bottleneck GPU
-        elif settings.model_inter_op_threads == 0 and self.providers == ["CPUExecutionProvider"]:
+        elif settings.model_inter_op_threads == 0 and self.provider == "CPUExecutionProvider":
             sess_options.inter_op_num_threads = 1
 
         # Set intra_op threads
         if settings.model_intra_op_threads > 0:
             sess_options.intra_op_num_threads = settings.model_intra_op_threads
-        elif settings.model_intra_op_threads == 0 and self.providers == ["CPUExecutionProvider"]:
+        elif settings.model_intra_op_threads == 0 and self.provider == "CPUExecutionProvider":
             sess_options.intra_op_num_threads = self.threads
 
         if sess_options.inter_op_num_threads > 1:
             sess_options.execution_mode = ort.ExecutionMode.ORT_PARALLEL
 
+        for dimension, size in self.overrides:
+            sess_options.add_free_dimension_override_by_name(dimension, size)
+
+        log.debug(f"Setting execution_mode to {sess_options.execution_mode.name}")
+        log.debug(f"Setting inter_op_num_threads to {sess_options.inter_op_num_threads}")
+        log.debug(f"Setting intra_op_num_threads to {sess_options.intra_op_num_threads}")
         return sess_options
 
 
+class Facts(BaseModel):
+    onnxruntime: str
+    options: dict[str, Any]
+    overrides: list[tuple[str, int]]
+    disabled_optimizers: list[str]
+    rewrite: str | None = None
+    device: str | None = None
+    half: bool = False
+
+
+class Manifest(BaseModel):
+    facts: Facts
+    graph: str  # relative to the artifact's directory
+
+
+def fresh(spec: GraphSpec) -> Path | None:
+    try:
+        manifest = Manifest.model_validate_json(spec.manifest.read_bytes())
+    except (OSError, ValueError):
+        return None
+    return spec.model_path.parent / manifest.graph if manifest.facts == spec.facts else None
+
+
+def prepared(spec: GraphSpec) -> Path:
+    if (graph := fresh(spec)) is None:
+        log.info(f"Preparing {spec.model_path} for {spec.provider} {dict(spec.pins)}")
+        # child process prepares the graph to avoid memory overhead in the main process
+        child = subprocess.run([sys.executable, "-m", "immich_ml.sessions.prepare"], input=pickle.dumps(spec))
+        if (graph := fresh(spec)) is None:
+            if child.returncode == UNREADABLE:
+                raise InvalidProtobuf(f"{spec.model_path} could not be read")
+            raise RuntimeError(f"Preparing {spec.model_path} failed; the output above says why")
+    return graph
+
+
+@cache
+def _plan(provider: str) -> RewritePlan:
+    version = tuple(int(piece) for piece in ort.__version__.split(".")[:3])
+    return plan_rewrites(RewriteContext(target=provider, ort_version=version))
+
+
+@cache
+def _cpu() -> str:
+    try:
+        cpuinfo = Path("/proc/cpuinfo").read_text()
+    except OSError:
+        return platform.machine()
+    return next(line for line in cpuinfo.splitlines() if line.startswith(("flags", "Features")))
+
+
+class OrtSession:
+    def __init__(
+        self, model_path: Path | str, policy: ShapePolicy, providers: list[str] | None = None, threads: int = 2
+    ) -> None:
+        self.model_path = Path(model_path)
+        self.policy = policy
+        self.providers = providers if providers is not None else _providers_default()
+        log.info(f"Setting execution providers to {self.providers}, in descending order of preference")
+        self.disabled_optimizers = _disabled_optimizers_default(self.providers)
+        log.debug(f"Setting disabled_optimizers to {self.disabled_optimizers}")
+        self.threads = threads
+        self.dynamic = self.providers[0] in DYNAMIC_PROVIDERS
+        # the shapes to snap onto, where a dim a shape leaves out takes any size
+        self.shapes = (policy.pinned,) if self.dynamic else policy.dims
+        self.batches = tuple(sorted({shape.batch for shape in self.shapes}, reverse=True))
+        self.graphs: dict[Shape, OrtGraph] = {}
+        self.locks: dict[Shape, Lock] = {}
+        if self.dynamic or len(policy.dims) == 1:
+            self.for_shape(policy.dims[0])  # the one graph there is, so nothing waits on a shape
+
+    def for_shape(self, shape: Shape) -> OrtGraph:
+        if self.dynamic:  # a batch gains nothing here, and a graph that knows it has one row optimizes further
+            shape = self.shapes[0]
+        if (graph := self.graphs.get(shape)) is not None:
+            return graph
+        if not self.dynamic and shape not in self.policy.dims:  # every shape admitted is another graph to compile
+            raise ValueError(f"Asked for {shape}, which is none of {list(self.policy.dims)}")
+        with self.locks.setdefault(shape, Lock()):  # requests arrive together for a graph none of them has
+            if (graph := self.graphs.get(shape)) is None:
+                log.debug(f"Building a graph for {shape}")
+                spec = GraphSpec(
+                    self.model_path,
+                    shape.pins,
+                    _overrides(self.policy, shape.pins),
+                    self.providers,
+                    self.disabled_optimizers,
+                    self.threads,
+                )
+                graph = self.graphs[shape] = OrtGraph(spec)
+            return graph
+
+    def warm(self) -> None:
+        """Opens a graph for every shape and runs it once before opening the next: OpenVINO defers work to the first
+        run of a graph it imports from its cache, which importing another graph first can corrupt."""
+        for shape in self.policy.dims[:1] if self.dynamic else self.policy.dims:
+            session = self.for_shape(shape).session
+            sizes = dict(_overrides(self.policy, shape.pins))  # for a dim the graph leaves free
+            feed = {
+                node.name: np.zeros(
+                    [dim if isinstance(dim, int) else sizes.get(dim, 1) for dim in node.shape], DTYPES[node.type]
+                )
+                for node in session.get_inputs()
+            }
+            session.run(None, feed)
+
+
+def _overrides(policy: ShapePolicy, pins: Mapping[str, int]) -> list[tuple[str, int]]:
+    order = tuple(policy.dims[0].pins)
+    named = []
+    for dimension, size in pins.items():
+        position = order.index(dimension)  # a graph that names several dims takes the first pin for all
+        named += [(dimension, size), (f"DynamicDimension.{position}", size)]
+        named.append(("None", size) if position == 0 else ("?", size))
+    return named
+
+
+@cache
+def _intel_gpu(device: str) -> Device:
+    """The IP version and execution units OpenVINO keys its blobs by, and the driver it checks them against."""
+    ov = ctypes.CDLL(str(Path(ort.__file__).parent / "capi" / "libopenvino_c.so"))
+    core, wanted = ctypes.c_void_p(), ctypes.c_char_p()
+    ov.ov_core_create(ctypes.byref(core))
+    ov.ov_core_get_property(core, device.encode(), b"DEVICE_UUID", ctypes.byref(wanted))
+    uuid = wanted.value
+    ov.ov_free(wanted)
+    ov.ov_core_free(core)
+    cl, gpu = ctypes.CDLL("libOpenCL.so.1"), ctypes.c_uint64(1 << 2)  # CL_DEVICE_TYPE_GPU
+    platforms, devices, count = (ctypes.c_void_p * 8)(), (ctypes.c_void_p * 8)(), ctypes.c_uint32()
+    cl.clGetPlatformIDs(8, platforms, ctypes.byref(count))
+    for vendor in platforms[: count.value]:
+        found = ctypes.c_uint32()
+        cl.clGetDeviceIDs(ctypes.c_void_p(vendor), gpu, 8, devices, ctypes.byref(found))
+        for handle in devices[: found.value]:
+            own, driver = (ctypes.c_ubyte * 16)(), ctypes.create_string_buffer(64)
+            ip, units = ctypes.c_uint32(), ctypes.c_uint32()
+            # CL_DEVICE_UUID_KHR, CL_DRIVER_VERSION, CL_DEVICE_IP_VERSION_INTEL, CL_DEVICE_MAX_COMPUTE_UNITS
+            for key, value in ((0x106A, own), (0x102D, driver), (0x4250, ip), (0x1002, units)):
+                cl.clGetDeviceInfo(ctypes.c_void_p(handle), key, ctypes.sizeof(value), ctypes.byref(value), None)
+            if bytes(own).hex().encode() == uuid:
+                version = f"{ip.value >> 22}.{ip.value >> 14 & 0xFF}.{ip.value & 0x3FFF}"
+                return Device(f"{version}-{units.value}eu", driver.value.decode())
+    raise LookupError(f"OpenCL has no GPU with the UUID of OpenVINO's {device}")
+
+
+@cache
+def _amd_gpu(index: int) -> Device:
+    """The arch MIGraphX keys its programs by, and the MIGraphX and HIP it loads, which its keys leave out."""
+    # the provider finds ROCm through its runpath
+    ctypes.CDLL(str(Path(ort.__file__).parent / "capi" / "libonnxruntime_providers_migraphx.so"))
+    with open("/proc/self/maps") as maps:
+        loaded = {Path(path).name.split(".so")[0]: Path(path) for *_, path in map(str.split, maps) if ".so" in path}
+    hip, hsa = ctypes.CDLL(str(loaded["libamdhip64"])), ctypes.CDLL(str(loaded["libhsa-runtime64"]))
+    wanted, arches = ctypes.create_string_buffer(64), {}
+    hip.hipDeviceGetPCIBusId(wanted, 64, index)  # honors HIP_VISIBLE_DEVICES, unlike HSA
+
+    def visit(agent: int, _: int | None) -> int:
+        name, bdf, domain = ctypes.create_string_buffer(64), ctypes.c_uint32(), ctypes.c_uint32()
+        # HSA_AGENT_INFO_NAME, HSA_AMD_AGENT_INFO_BDFID, HSA_AMD_AGENT_INFO_DOMAIN
+        for key, value in ((0, name), (0xA006, bdf), (0xA00F, domain)):
+            hsa.hsa_agent_get_info(ctypes.c_uint64(agent), key, ctypes.byref(value))
+        bus, device, function = bdf.value >> 8 & 0xFF, bdf.value >> 3 & 0x1F, bdf.value & 0x7
+        arches[f"{domain.value:04x}:{bus:02x}:{device:02x}.{function:01x}"] = name.value.decode()
+        return 0
+
+    hsa.hsa_init()
+    hsa.hsa_iterate_agents(ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_uint64, ctypes.c_void_p)(visit), None)
+    return Device(arches[wanted.value.decode()], f"{loaded['libmigraphx'].name} {loaded['libamdhip64'].name}")
+
+
+def flush_denormals() -> None:
+    """Reads subnormal floats as zero on the calling thread, which ORT does only on its own threads."""
+    if sys.platform != "linux" or platform.machine() != "x86_64":
+        return
+    libm = ctypes.CDLL("libm.so.6")
+    env = (ctypes.c_uint32 * 8)()  # glibc's fenv_t here: the x87 environment, then MXCSR
+    libm.fegetenv(env)
+    env[7] |= 0x8040  # denormals are zero, flush to zero
+    libm.fesetenv(env)
+
+
+def _providers_default() -> list[str]:
+    from immich_ml.models.constants import SUPPORTED_PROVIDERS  # here, as the preparing child has no use for the models
+
+    available_providers = set(ort.get_available_providers())
+    log.debug(f"Available ORT providers: {available_providers}")
+    return [provider for provider in SUPPORTED_PROVIDERS if provider in available_providers]
+
+
 def _disabled_optimizers_default(providers: list[str]) -> list[str]:
     disabled_optimizers: list[str] = []
     if platform.machine() in ("arm64", "aarch64"):  # as macOS and Linux name the same architecture
@@ -226,14 +409,3 @@ def _disabled_optimizers_default(providers: list[str]) -> list[str]:
         disabled_optimizers.append("MatMulAddFusion")
 
     return disabled_optimizers
-
-
-def flush_denormals() -> None:
-    """Reads subnormal floats as zero on the calling thread, which ORT does only on its own threads."""
-    if sys.platform != "linux" or platform.machine() != "x86_64":
-        return
-    libm = ctypes.CDLL("libm.so.6")
-    env = (ctypes.c_uint32 * 8)()  # glibc's fenv_t here: the x87 environment, then MXCSR
-    libm.fegetenv(env)
-    env[7] |= 0x8040  # denormals are zero, flush to zero
-    libm.fesetenv(env)

machine-learning/immich_ml/sessions/policy.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+from immich_ml.schemas import Shape
+
+
+@dataclass(frozen=True)
+class ShapePolicy:
+    dims: tuple[Shape, ...] = (Shape(batch=1),)
+    label: str = ""  # the artifact subdirectory, for a submodel that ships more than one set
+
+    @property
+    def pinned(self) -> Shape:
+        first = self.dims[0]
+        return Shape(
+            batch=1,
+            height=first.height if all(shape.height == first.height for shape in self.dims) else None,
+            width=first.width if all(shape.width == first.width for shape in self.dims) else None,
+        )
+
+
+def batches(maximum: int) -> tuple[int, ...]:
+    return (1, maximum) if maximum > 1 else (1,)
+
+
+def runs(total: int, batches: Sequence[int]) -> list[int]:
+    runs = []
+    while total > 0:
+        runs.append(max(size for size in batches if size <= total))
+        total -= runs[-1]
+    return runs

machine-learning/immich_ml/sessions/prepare.py

@@ -0,0 +1,107 @@
+import fcntl
+import pickle
+import sys
+from pathlib import Path
+from shutil import rmtree
+from typing import Any
+
+import onnx
+from google.protobuf.message import DecodeError
+from immich_model.onnx.f16 import derive
+from immich_model.runtime import apply_rewrites
+from onnxruntime.capi.onnxruntime_pybind11_state import InvalidProtobuf, NoSuchFile
+
+from ..config import settings
+from .ort import UNREADABLE, GraphSpec, Manifest, fresh
+
+
+def main() -> None:
+    spec: GraphSpec = pickle.loads(sys.stdin.buffer.read())
+    with spec.model_path.with_suffix(".lock").open("w") as lock:
+        fcntl.flock(lock, fcntl.LOCK_EX)  # several workers may want the same graph
+        if fresh(spec) is None:
+            try:
+                prepare(spec)
+            except (InvalidProtobuf, NoSuchFile, DecodeError):
+                sys.exit(UNREADABLE)
+
+
+def prepare(spec: GraphSpec) -> None:
+    rmtree(spec.directory, ignore_errors=True)  # whatever is here was made under other facts, or not finished
+    spec.directory.mkdir(parents=True)
+    graph = rewritten(spec, narrowed(spec))
+    if spec.provider == "CPUExecutionProvider":
+        graph = prepacked(spec, graph)  # allows mmap'ing the file instead of prepacking to dirty heap memory
+    elif spec.directory.as_posix() in spec.provider_options[0].values():
+        spec.session(graph)  # the provider keeps what it compiles, in the directory it was given
+    manifest = Manifest(facts=spec.facts, graph=graph.relative_to(spec.model_path.parent).as_posix())
+    spec.manifest.write_text(manifest.model_dump_json())  # last, so that it means complete
+
+
+def narrowed(spec: GraphSpec) -> Path:
+    if not spec.half:
+        return spec.model_path
+    destination = spec.directory / "model_fp16.onnx"
+    derive(spec.model_path, destination)
+    return destination
+
+
+def rewritten(spec: GraphSpec, graph: Path) -> Path:
+    written: Path = graph if settings.legacy_models else apply_rewrites(graph, spec.plan)
+    return written
+
+
+def prepacked(spec: GraphSpec, graph: Path) -> Path:
+    destination = spec.directory / "model.onnx"
+    options = spec.sess_options()
+    options.optimized_model_filepath = destination.as_posix()
+    options.add_session_config_entry("session.optimized_model_external_initializers_file_name", "model.data")
+    options.add_session_config_entry("session.save_external_prepacked_constant_initializers", "1")
+    options.add_session_config_entry("session.optimized_model_external_initializers_min_size_in_bytes", "1024")
+    spec.session(graph, options)
+    align(destination)
+    return destination
+
+
+def align(destination: Path) -> None:
+    """Re-lays the data file so every region starts aligned. ORT pads only large regions, so one odd length skews
+    all that follow, and MLAS reads a packed weight with aligned loads: it faults, and reads the rest slower."""
+    model = onnx.load(destination.as_posix(), load_external_data=False)
+    length_at: dict[int, int] = {}
+    offsets: list[Any] = []
+    packs: list[tuple[Any, str, list[list[str]]]] = []
+    for tensor in model.graph.initializer:
+        fields = {entry.key: entry for entry in tensor.external_data}
+        for key, entry in fields.items():
+            if key == "offset":
+                length_at[int(entry.value)] = int(fields["length"].value)
+                offsets.append(entry)
+            elif key.startswith("prepacked_"):
+                blob_key, *rest = entry.value.split("|")
+                triples = [item.split(";") for item in rest]
+                length_at.update((int(offset), int(length)) for offset, length, _ in triples)
+                packs.append((entry, blob_key, triples))
+
+    moved, cursor = {}, 0
+    for offset in sorted(length_at):
+        cursor += -cursor % (4096 if length_at[offset] > 1 << 20 else 64)
+        moved[offset] = cursor
+        cursor += length_at[offset]
+
+    data = destination.with_suffix(".data")
+    with data.open("rb") as source, data.with_suffix(".aligned").open("wb") as target:
+        for offset in sorted(length_at):
+            source.seek(offset)
+            target.seek(moved[offset])  # what it skips reads as zeros
+            target.write(source.read(length_at[offset]))
+    data.with_suffix(".aligned").replace(data)
+
+    for entry in offsets:
+        entry.value = str(moved[int(entry.value)])
+    for entry, blob_key, triples in packs:
+        entry.value = "|".join([blob_key, *(f"{moved[int(o)]};{n};{c}" for o, n, c in triples)])
+    onnx.save(model, destination.as_posix())
+
+
+if __name__ == "__main__":
+    main()

machine-learning/immich_ml/sessions/rknn/__init__.py

@@ -1,78 +1,84 @@
 from __future__ import annotations
 
+from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, NamedTuple
+from typing import Any
 
 import numpy as np
+import orjson
 from numpy.typing import NDArray
 
 from immich_ml.config import log, settings
-from immich_ml.schemas import SessionNode
+from immich_ml.schemas import ModelInput, ModelTensor, SessionNode, Shape
 
-from .rknnpool import RknnPoolExecutor, is_available, soc_name
+from .rknnpool import RknnNode, RknnPoolExecutor, is_available, native_outputs, soc_name
 
 is_available = is_available and settings.rknn
 model_prefix = Path("rknpu") / soc_name if is_available and soc_name is not None else None
 
 
-def run_inference(rknn_lite: Any, input: list[NDArray[np.float32]]) -> list[NDArray[np.float32]]:
-    outputs: list[NDArray[np.float32]] = rknn_lite.inference(inputs=input, data_format="nchw")
-    return outputs
+def model_path(model_dir: Path, variant: str = "") -> Path:
+    return (model_dir / model_prefix if model_prefix else model_dir) / variant / "model.rknn"
 
 
-input_output_mapping: dict[str, dict[str, Any]] = {
-    "detection": {
-        "input": {"norm_tensor:0": (1, 3, 640, 640)},
-        "output": {
-            "norm_tensor:1": (12800, 1),
-            "norm_tensor:2": (3200, 1),
-            "norm_tensor:3": (800, 1),
-            "norm_tensor:4": (12800, 4),
-            "norm_tensor:5": (3200, 4),
-            "norm_tensor:6": (800, 4),
-            "norm_tensor:7": (12800, 10),
-            "norm_tensor:8": (3200, 10),
-            "norm_tensor:9": (800, 10),
-        },
-    },
-    "recognition": {"input": {"norm_tensor:0": (1, 3, 112, 112)}, "output": {"norm_tensor:1": (1, 512)}},
-}
+def run_inference(rknn_lite: Any, inputs: list[ModelTensor], data_format: str | None) -> list[NDArray[np.float32]]:
+    rknn_lite.rknn_runtime.set_inputs(inputs, None, data_format)
+    rknn_lite.rknn_runtime.run(False)
+    return native_outputs(rknn_lite)
+
+
+def input_layout(compiled: tuple[int, ...], array: ModelTensor) -> str | None:
+    if array.ndim != 4:
+        return None  # token ids carry no layout
+    if array.shape[3] == compiled[3]:
+        return "nhwc"
+    if array.shape[1] == compiled[3]:
+        return "nchw"
+    raise ValueError(f"binary takes {compiled[3]} channels, fed {list(array.shape)}")
 
 
 class RknnSession:
     def __init__(self, model_path: Path) -> None:
-        self.model_type = "detection" if "detection" in model_path.parts else "recognition"
         self.tpe = settings.rknn_threads
 
         log.info(f"Loading RKNN model from {model_path} with {self.tpe} threads.")
         self.rknnpool = RknnPoolExecutor(model_path=model_path.as_posix(), tpes=self.tpe, func=run_inference)
         log.info(f"Loaded RKNN model from {model_path} with {self.tpe} threads.")
+        # the shapes the binary was compiled for, among which the runtime routes itself
+        batch = self.rknnpool.inputs[0].shape[0]
+        declared = orjson.loads(self.rknnpool.custom_string) if self.rknnpool.custom_string else {}
+        self.shapes = tuple(Shape(batch, **dims) for dims in declared.get("dims", [{}]))
+        self.metadata = {key: value for key, value in declared.items() if key != "dims"}
+        self.batches = (batch,)
 
-    def get_inputs(self) -> list[SessionNode]:
-        return [RknnNode(name=k, shape=v) for k, v in input_output_mapping[self.model_type]["input"].items()]
+    def for_shape(self, shape: Shape) -> RknnSession:
+        return self
 
-    def get_outputs(self) -> list[SessionNode]:
-        return [RknnNode(name=k, shape=v) for k, v in input_output_mapping[self.model_type]["output"].items()]
+    def warm(self) -> None:
+        pass  # the binary is loaded whole when the session opens
+
+    def get_inputs(self) -> Sequence[SessionNode]:
+        return self.rknnpool.inputs
+
+    def get_outputs(self) -> Sequence[SessionNode]:
+        return self.rknnpool.outputs
 
     def get_metadata(self) -> dict[str, str]:
-        return {}
+        return self.metadata
+
+    @property
+    def normalizes_input(self) -> bool:
+        # the compiler transposes and casts every binary, so only what the exporter declared can say
+        return bool(self.rknnpool.custom_string)
 
     def run(
         self,
         output_names: list[str] | None,
-        input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]] | dict[str, NDArray[np.uint8]],
+        input_feed: ModelInput,
         run_options: Any = None,
     ) -> list[NDArray[np.float32]]:
-        input_data: list[NDArray[np.float32]] = [np.ascontiguousarray(v) for v in input_feed.values()]
-        res = self.rknnpool.run(input_data)
-        if res is None:
-            raise RuntimeError("RKNN inference failed!")
-        return res
+        inputs = [array if array.flags.c_contiguous else array.copy() for array in input_feed.values()]
+        return self.rknnpool.run(inputs, input_layout(self.rknnpool.inputs[0].shape, inputs[0]))
 
 
-class RknnNode(NamedTuple):
-    name: str
-    shape: tuple[int, ...]
-
-
-__all__ = ["RknnSession", "RknnNode", "is_available", "soc_name", "model_prefix"]
+__all__ = ["RknnSession", "RknnNode", "is_available", "soc_name", "model_path"]

machine-learning/immich_ml/sessions/rknn/rknnpool.py

@@ -1,17 +1,19 @@
 # This code is from leafqycc/rknn-multi-threaded
 # Following Apache License 2.0
 
+import ctypes
 import logging
 import threading
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
-from typing import Callable
+from typing import Any, Callable, NamedTuple
 
 import numpy as np
 from numpy.typing import NDArray
 
 from immich_ml.config import log
 from immich_ml.models.constants import RKNN_COREMASK_SUPPORTED_SOCS, RKNN_SUPPORTED_SOCS
+from immich_ml.schemas import ModelTensor
 
 
 def get_soc(device_tree_path: Path | str) -> str | None:
@@ -30,7 +32,7 @@ def get_soc(device_tree_path: Path | str) -> str | None:
 soc_name = None
 is_available = False
 try:
-    from rknnlite.api import RKNNLite
+    from rknnlite.api import RKNNLite, rknn_runtime
 
     soc_name = get_soc("/proc/device-tree/compatible")
     is_available = soc_name is not None
@@ -38,6 +40,82 @@ except ImportError:
     log.debug("RKNN is not available")
 
 
+class RknnNode(NamedTuple):
+    name: str
+    shape: tuple[int, ...]
+
+
+_CUSTOM_STRING_QUERY = 7  # RKNN_QUERY_CUSTOM_STRING
+# RKNN_QUERY_OUTPUT_ATTR, and RKNN_QUERY_CURRENT_OUTPUT_ATTR for the shape a dynamic binary last produced
+_OUTPUT_QUERY, _CURRENT_OUTPUT_QUERY = 2, 15
+_DTYPES = {0: np.float32, 1: np.float16}  # rknn_tensor_type
+
+
+class _CustomString(ctypes.Structure):
+    _fields_ = [("string", ctypes.c_char * 1024)]
+
+
+class _TensorAttr(ctypes.Structure):  # rknn_tensor_attr
+    _fields_ = [
+        ("index", ctypes.c_uint32),
+        ("n_dims", ctypes.c_uint32),
+        ("dims", ctypes.c_uint32 * 16),
+        ("name", ctypes.c_char * 256),
+        ("n_elems", ctypes.c_uint32),
+        ("size", ctypes.c_uint32),
+        ("fmt", ctypes.c_int),
+        ("type", ctypes.c_int),
+        ("qnt_type", ctypes.c_int),
+        ("fl", ctypes.c_int8),
+        ("zp", ctypes.c_int32),
+        ("scale", ctypes.c_float),
+        ("w_stride", ctypes.c_uint32),
+        ("size_with_stride", ctypes.c_uint32),
+        ("pass_through", ctypes.c_uint8),
+        ("h_stride", ctypes.c_uint32),
+    ]
+
+
+def native_outputs(rknn_lite: "RKNNLite") -> list[NDArray[Any]]:
+    """The outputs in the precision the NPU computed them in, which RKNNLite would widen to fp32 on the CPU."""
+    runtime: Any = rknn_lite.rknn_runtime
+    count = runtime.get_in_out_num()[1]
+    buffers = (rknn_runtime.RKNNOutput * count)()
+    for index, buffer in enumerate(buffers):
+        buffer.index = index
+    if runtime.lib.rknn_outputs_get(runtime.context, count, buffers, None) != 0:
+        raise RuntimeError("RKNN inference failed!")
+    query = _CURRENT_OUTPUT_QUERY if runtime.is_dynamic_shape else _OUTPUT_QUERY
+    outputs = []
+    for buffer in buffers:
+        attr = _TensorAttr(index=buffer.index)
+        runtime.lib.rknn_query(runtime.context, query, ctypes.byref(attr), ctypes.sizeof(attr))
+        data = ctypes.cast(buffer.buf, ctypes.POINTER(ctypes.c_char * buffer.size)).contents
+        outputs.append(np.frombuffer(data, _DTYPES[attr.type]).reshape(attr.dims[: attr.n_dims]).copy())
+    runtime.lib.rknn_outputs_release(runtime.context, count, buffers)
+    return outputs
+
+
+def tensor_nodes(rknn_lite: "RKNNLite") -> tuple[list[RknnNode], list[RknnNode]]:
+    runtime = rknn_lite.rknn_runtime
+    n_inputs, n_outputs = runtime.get_in_out_num()
+
+    def node(index: int, is_output: bool) -> RknnNode:
+        attr = runtime.get_tensor_attr(index, is_output=is_output)
+        return RknnNode(attr.name.decode(), tuple(attr.dims[: attr.n_dims]))
+
+    return [node(i, False) for i in range(n_inputs)], [node(i, True) for i in range(n_outputs)]
+
+
+def custom_string(rknn_lite: "RKNNLite") -> str:
+    runtime: Any = rknn_lite.rknn_runtime
+    buffer = _CustomString()
+    # rknnlite declares argtypes on the shared handle, so pass plain values and let ctypes convert
+    runtime.lib.rknn_query(runtime.context, _CUSTOM_STRING_QUERY, ctypes.byref(buffer), ctypes.sizeof(buffer))
+    declared: str = buffer.string.decode()
+    return declared
+
+
 def init_rknn(model_path: str) -> "RKNNLite":
     if not is_available:
         raise RuntimeError("rknn is not available!")
@@ -63,21 +141,23 @@ class RknnPoolExecutor:
         self,
         model_path: str,
         tpes: int,
-        func: Callable[["RKNNLite", list[NDArray[np.float32]]], list[NDArray[np.float32]]],
+        func: Callable[["RKNNLite", list[ModelTensor], str | None], list[NDArray[np.float32]]],
     ) -> None:
         self.tpes = tpes
         self.rknn_pool = [init_rknn(model_path) for _ in range(tpes)]
+        self.inputs, self.outputs = tensor_nodes(self.rknn_pool[0])
+        self.custom_string = custom_string(self.rknn_pool[0])
         self.pool = ThreadPoolExecutor(max_workers=tpes)
         self.func = func
         self.num = 0
         self.lock = threading.Lock()
 
-    def run(self, inputs: list[NDArray[np.float32]]) -> list[NDArray[np.float32]]:
+    def run(self, inputs: list[ModelTensor], data_format: str | None) -> list[NDArray[np.float32]]:
         with self.lock:
             idx = self.num % self.tpes
             self.num += 1
 
-        fut = self.pool.submit(self.func, self.rknn_pool[idx], inputs)
+        fut = self.pool.submit(self.func, self.rknn_pool[idx], inputs, data_format)
         return fut.result()
 
     def release(self) -> None: