feat(ml): update ONNX Runtime, OpenVINO and ROCm stack

.editorconfig

@@ -8,6 +8,9 @@ indent_size = 2
 insert_final_newline = true
 trim_trailing_whitespace = true
 
+[*.py]
+indent_size = 4
+
 [*.{ts,js}]
 quote_type = single
 

.github/workflows/test.yml

@@ -565,9 +565,9 @@ jobs:
         uses: astral-sh/setup-uv@2ddd2b9cb38ad8efd50337e8ab201519a34c9f24 # v7.1.1
       - uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c # v6.0.0
         # TODO: add caching when supported (https://github.com/actions/setup-python/pull/818)
-        # with:
-        #   python-version: 3.11
-        #   cache: 'uv'
+        with:
+          python-version: 3.11
+          #cache: 'uv'
       - name: Install dependencies
         run: |
           uv sync --extra cpu

docs/docs/install/environment-variables.md

@@ -172,6 +172,7 @@ Redis (Sentinel) URL example JSON before encoding:
 | `MACHINE_LEARNING_RKNN`                                     | Enable RKNN hardware acceleration if supported                                                      |             `True`              | machine learning |
 | `MACHINE_LEARNING_RKNN_THREADS`                             | How many threads of RKNN runtime should be spinned up while inferencing.                            |               `1`               | machine learning |
 | `MACHINE_LEARNING_MODEL_ARENA`                              | Pre-allocates CPU memory to avoid memory fragmentation                                              |              true               | machine learning |
+| `MACHINE_LEARNING_OPENVINO_CACHE_CAPACITY`                  | The max number of image dimensions for which models have cached optimizations                       |              `20`               | machine learning |
 
 \*1: It is recommended to begin with this parameter when changing the concurrency levels of the machine learning service and then tune the other ones.
 

machine-learning/Dockerfile

@@ -2,7 +2,7 @@ ARG DEVICE=cpu
 
 FROM python:3.11-bookworm@sha256:fc1f2e357c307c4044133952b203e66a47e7726821a664f603a180a0c5823844 AS builder-cpu
 
-FROM builder-cpu AS builder-openvino
+FROM python:3.13-slim-trixie@sha256:0222b795db95bf7412cede36ab46a266cfb31f632e64051aac9806dabf840a61 AS builder-openvino
 
 FROM builder-cpu AS builder-cuda
 
@@ -22,20 +22,20 @@ FROM builder-cpu AS builder-rknn
 
 # Warning: 25GiB+ disk space required to pull this image
 # TODO: find a way to reduce the image size
-FROM rocm/dev-ubuntu-22.04:6.4.3-complete@sha256:1f7e92ca7e3a3785680473329ed1091fc99db3e90fcb3a1688f2933e870ed76b AS builder-rocm
+FROM rocm/dev-ubuntu-24.04:6.4.4-complete@sha256:31418ac10a3769a71eaef330c07280d1d999d7074621339b8f93c484c35f6078 AS builder-rocm
 
 # renovate: datasource=github-releases depName=Microsoft/onnxruntime
-ARG ONNXRUNTIME_VERSION="v1.20.1"
+ARG ONNXRUNTIME_VERSION="v1.21.1"
 WORKDIR /code
 
-RUN apt-get update && apt-get install -y --no-install-recommends wget git python3.10-venv
-RUN wget -nv https://github.com/Kitware/CMake/releases/download/v3.30.1/cmake-3.30.1-linux-x86_64.sh && \
-    chmod +x cmake-3.30.1-linux-x86_64.sh && \
-    mkdir -p /code/cmake-3.30.1-linux-x86_64 && \
-    ./cmake-3.30.1-linux-x86_64.sh --skip-license --prefix=/code/cmake-3.30.1-linux-x86_64 && \
-    rm cmake-3.30.1-linux-x86_64.sh
+RUN apt-get update && apt-get install -y --no-install-recommends wget git python3.12-venv
+RUN wget -nv https://github.com/Kitware/CMake/releases/download/v3.31.9/cmake-3.31.9-linux-x86_64.sh && \
+    chmod +x cmake-3.31.9-linux-x86_64.sh && \
+    mkdir -p /code/cmake-3.31.9-linux-x86_64 && \
+    ./cmake-3.31.9-linux-x86_64.sh --skip-license --prefix=/code/cmake-3.31.9-linux-x86_64 && \
+    rm cmake-3.31.9-linux-x86_64.sh
 
-ENV PATH=/code/cmake-3.30.1-linux-x86_64/bin:${PATH}
+ENV PATH=/code/cmake-3.31.9-linux-x86_64/bin:${PATH}
 
 RUN git clone --single-branch --branch "${ONNXRUNTIME_VERSION}" --recursive "https://github.com/Microsoft/onnxruntime" onnxruntime
 WORKDIR /code/onnxruntime
@@ -45,9 +45,25 @@ COPY ./patches/* /tmp/
 RUN git apply /tmp/*.patch
 
 RUN /bin/sh ./dockerfiles/scripts/install_common_deps.sh
+
+ENV CCACHE_DIR="/ccache"
 # Note: the `parallel` setting uses a substantial amount of RAM
-RUN ./build.sh --allow_running_as_root --config Release --build_wheel --update --build --parallel 17 --cmake_extra_defines\
-    ONNXRUNTIME_VERSION="${ONNXRUNTIME_VERSION}" --skip_tests --use_rocm --rocm_home=/opt/rocm
+RUN --mount=type=cache,target=/ccache \
+    ./build.sh \
+    --allow_running_as_root \
+    --config Release \
+    --build_wheel \
+    --update \
+    --build \
+    --parallel 17 \
+    --cmake_extra_defines \
+    ONNXRUNTIME_VERSION="${ONNXRUNTIME_VERSION}" \
+    CMAKE_HIP_ARCHITECTURES="gfx900;gfx908;gfx90a;gfx1030;gfx1100;gfx1101;gfx1102;gfx940;gfx941;gfx942;gfx1200;gfx1201" \
+    --skip_tests \
+    --use_rocm \
+    --rocm_home=/opt/rocm \
+    --use_cache \
+    --compile_no_warning_as_error
 RUN mv /code/onnxruntime/build/Linux/Release/dist/*.whl /opt/
 
 FROM builder-${DEVICE} AS builder
@@ -73,15 +89,18 @@ FROM python:3.11-slim-bookworm@sha256:873f91540d53b36327ed4fb018c9669107a4e2a676
 ENV LD_PRELOAD=/usr/lib/libmimalloc.so.2 \
     MACHINE_LEARNING_MODEL_ARENA=false
 
-FROM python:3.11-slim-bookworm@sha256:873f91540d53b36327ed4fb018c9669107a4e2a676719720edb4209c4b15d029 AS prod-openvino
+FROM python:3.13-slim-trixie@sha256:0222b795db95bf7412cede36ab46a266cfb31f632e64051aac9806dabf840a61 AS prod-openvino
 
 RUN apt-get update && \
     apt-get install --no-install-recommends -yqq ocl-icd-libopencl1 wget && \
-    wget -nv https://github.com/intel/intel-graphics-compiler/releases/download/igc-1.0.17384.11/intel-igc-core_1.0.17384.11_amd64.deb && \
-    wget -nv https://github.com/intel/intel-graphics-compiler/releases/download/igc-1.0.17384.11/intel-igc-opencl_1.0.17384.11_amd64.deb && \
-    wget -nv https://github.com/intel/compute-runtime/releases/download/24.31.30508.7/intel-opencl-icd_24.31.30508.7_amd64.deb && \
+    wget -nv https://github.com/intel/intel-graphics-compiler/releases/download/v2.20.3/intel-igc-core-2_2.20.3+19972_amd64.deb && \
+    wget -nv https://github.com/intel/intel-graphics-compiler/releases/download/v2.20.3/intel-igc-opencl-2_2.20.3+19972_amd64.deb && \
+    wget -nv https://github.com/intel/compute-runtime/releases/download/25.40.35563.4/intel-opencl-icd_25.40.35563.4-0_amd64.deb &&  \
+    wget -nv https://github.com/intel/intel-graphics-compiler/releases/download/igc-1.0.17537.24/intel-igc-core_1.0.17537.24_amd64.deb && \
+    wget -nv https://github.com/intel/intel-graphics-compiler/releases/download/igc-1.0.17537.24/intel-igc-opencl_1.0.17537.24_amd64.deb && \
+    wget -nv https://github.com/intel/compute-runtime/releases/download/24.35.30872.36/intel-opencl-icd-legacy1_24.35.30872.36_amd64.deb && \
     # TODO: Figure out how to get renovate to manage this differently versioned libigdgmm file
-    wget -nv https://github.com/intel/compute-runtime/releases/download/24.31.30508.7/libigdgmm12_22.4.1_amd64.deb && \
+    wget -nv https://github.com/intel/compute-runtime/releases/download/25.40.35563.4/libigdgmm12_22.8.2_amd64.deb && \
     dpkg -i *.deb && \
     rm *.deb && \
     apt-get remove wget -yqq && \
@@ -102,7 +121,7 @@ COPY --from=builder-cuda /usr/local/bin/python3 /usr/local/bin/python3
 COPY --from=builder-cuda /usr/local/lib/python3.11 /usr/local/lib/python3.11
 COPY --from=builder-cuda /usr/local/lib/libpython3.11.so /usr/local/lib/libpython3.11.so
 
-FROM rocm/dev-ubuntu-22.04:6.4.3-complete@sha256:1f7e92ca7e3a3785680473329ed1091fc99db3e90fcb3a1688f2933e870ed76b AS prod-rocm
+FROM rocm/dev-ubuntu-24.04:6.4.4-complete@sha256:31418ac10a3769a71eaef330c07280d1d999d7074621339b8f93c484c35f6078 AS prod-rocm
 
 FROM prod-cpu AS prod-armnn
 

machine-learning/immich_ml/config.py

@@ -70,6 +70,7 @@ class Settings(BaseSettings):
     rknn_threads: int = 1
     preload: PreloadModelData | None = None
     max_batch_size: MaxBatchSize | None = None
+    openvino_cache_capacity: int = 20
 
     @property
     def device_id(self) -> str:

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

@@ -3,10 +3,10 @@
 import cv2
 import numpy as np
 from numpy.typing import NDArray
-from PIL import Image
+from PIL import Image, ImageOps
 from rapidocr.ch_ppocr_det.utils import DBPostProcess
 from rapidocr.inference_engine.base import FileInfo, InferSession
-from rapidocr.utils import DownloadFile, DownloadFileInput
+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
 
@@ -78,18 +78,15 @@ def _predict(self, inputs: Image.Image) -> TextDetectionOutput:
 
     # adapted from RapidOCR
     def _transform(self, img: Image.Image) -> NDArray[np.float32]:
-        if img.height < img.width:
-            ratio = float(self.max_resolution) / img.height
+        aspect_ratio = img.width / img.height
+        if aspect_ratio > 1.25:
+            target_dims = (self.max_resolution * 2, self.max_resolution)
+        elif aspect_ratio < 0.75:
+            target_dims = (self.max_resolution, self.max_resolution * 2)
         else:
-            ratio = float(self.max_resolution) / img.width
-
-        resize_h = int(img.height * ratio)
-        resize_w = int(img.width * ratio)
-
-        resize_h = int(round(resize_h / 32) * 32)
-        resize_w = int(round(resize_w / 32) * 32)
-        resized_img = img.resize((int(resize_w), int(resize_h)), resample=Image.Resampling.LANCZOS)
+            target_dims = (self.max_resolution, self.max_resolution)
 
+        resized_img = ImageOps.pad(img, target_dims, color=(0, 0, 0), method=Image.Resampling.LANCZOS)
         img_np: NDArray[np.float32] = cv2.cvtColor(np.array(resized_img, dtype=np.float32), cv2.COLOR_RGB2BGR)  # type: ignore
         img_np -= self.mean
         img_np *= self.std_inv
@@ -116,8 +113,8 @@ def sorted_boxes(self, dt_boxes: NDArray[np.float32]) -> NDArray[np.float32]:
         return sorted_boxes
 
     def configure(self, **kwargs: Any) -> None:
-        if (max_resolution := kwargs.get("maxResolution")) is not None:
-            self.max_resolution = max_resolution
+        if (max_resolution := kwargs.get("maxResolution")) is not None and max_resolution != self.max_resolution:
+            self.max_resolution = int(round(max_resolution / 32) * 32)
         if (min_score := kwargs.get("minScore")) is not None:
             self.postprocess.box_thresh = min_score
         if (score_mode := kwargs.get("scoreMode")) is not None:

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

@@ -6,7 +6,7 @@
 from rapidocr.ch_ppocr_rec import TextRecInput
 from rapidocr.ch_ppocr_rec import TextRecognizer as RapidTextRecognizer
 from rapidocr.inference_engine.base import FileInfo, InferSession
-from rapidocr.utils import DownloadFile, DownloadFileInput
+from rapidocr.utils.download_file import DownloadFile, DownloadFileInput
 from rapidocr.utils.typings import EngineType, LangRec, OCRVersion, TaskType
 from rapidocr.utils.typings import ModelType as RapidModelType
 from rapidocr.utils.vis_res import VisRes

machine-learning/immich_ml/sessions/ort.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 import onnxruntime as ort
+import orjson
 from numpy.typing import NDArray
 
 from immich_ml.models.constants import SUPPORTED_PROVIDERS
@@ -97,6 +98,13 @@ def _provider_options_default(self) -> list[dict[str, Any]]:
                         "device_type": f"GPU.{settings.device_id}",
                         "precision": "FP32",
                         "cache_dir": (self.model_path.parent / "openvino").as_posix(),
+                        "load_config": orjson.dumps(
+                            {
+                                f"GPU.{settings.device_id}": {
+                                    "CPU_RUNTIME_CACHE_CAPACITY": str(settings.openvino_cache_capacity)
+                                },
+                            }
+                        ).decode(),
                     }
                 case "CoreMLExecutionProvider":
                     options = {

machine-learning/patches/0002-install-system-deps.patch

@@ -0,0 +1,26 @@
+diff --git a/dockerfiles/scripts/install_common_deps.sh b/dockerfiles/scripts/install_common_deps.sh
+index bbb672a99e..8c6430d919 100644
+--- a/dockerfiles/scripts/install_common_deps.sh
++++ b/dockerfiles/scripts/install_common_deps.sh
+@@ -8,16 +8,17 @@ apt-get update && apt-get install -y --no-install-recommends \
+         curl \
+         libcurl4-openssl-dev \
+         libssl-dev \
+-        python3-dev
++        python3-dev \
++        python3-numpy \
++        python3-packaging \
++        python3-wheel \
++        ccache
+
+ # Dependencies: conda
+ wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-4.5.11-Linux-x86_64.sh -O ~/miniconda.sh --no-check-certificate && /bin/bash ~/miniconda.sh -b -p /opt/miniconda
+ rm ~/miniconda.sh
+ /opt/miniconda/bin/conda clean -ya
+
+-pip install numpy
+-pip install packaging
+-pip install "wheel>=0.35.1"
+ rm -rf /opt/miniconda/pkgs
+
+ # Dependencies: cmake

machine-learning/patches/0003-eigen-switch-to-mirror-backport.patch

@@ -0,0 +1,15 @@
+diff --git a/cmake/deps.txt b/cmake/deps.txt
+index d0bab93d3c..1edc309ed0 100644
+--- a/cmake/deps.txt
++++ b/cmake/deps.txt
+@@ -22,7 +22,9 @@ dlpack;https://github.com/dmlc/dlpack/archive/refs/tags/v0.6.zip;4d565dd2e5b3132
+ # it contains changes on top of 3.4.0 which are required to fix build issues.
+ # Until the 3.4.1 release this is the best option we have.
+ # Issue link: https://gitlab.com/libeigen/eigen/-/issues/2744
+-eigen;https://gitlab.com/libeigen/eigen/-/archive/1d8b82b0740839c0de7f1242a3585e3390ff5f33/eigen-1d8b82b0740839c0de7f1242a3585e3390ff5f33.zip;5ea4d05e62d7f954a46b3213f9b2535bdd866803
++# Moved to github mirror to avoid gitlab issues.
++# Issue link: https://github.com/bazelbuild/bazel-central-registry/issues/4355
++eigen;https://github.com/eigen-mirror/eigen/archive/1d8b82b0740839c0de7f1242a3585e3390ff5f33/eigen-1d8b82b0740839c0de7f1242a3585e3390ff5f33.zip;05b19b49e6fbb91246be711d801160528c135e34
+ flatbuffers;https://github.com/google/flatbuffers/archive/refs/tags/v23.5.26.zip;59422c3b5e573dd192fead2834d25951f1c1670c
+ fp16;https://github.com/Maratyszcza/FP16/archive/0a92994d729ff76a58f692d3028ca1b64b145d91.zip;b985f6985a05a1c03ff1bb71190f66d8f98a1494
+ fxdiv;https://github.com/Maratyszcza/FXdiv/archive/63058eff77e11aa15bf531df5dd34395ec3017c8.zip;a5658f4036402dbca7cebee32be57fb8149811e1

machine-learning/pyproject.toml

@@ -3,7 +3,7 @@ name = "immich-ml"
 version = "2.2.3"
 description = ""
 authors = [{ name = "Hau Tran", email = "alex.tran1502@gmail.com" }]
-requires-python = ">=3.10,<4.0"
+requires-python = ">=3.11,<4"
 readme = "README.md"
 dependencies = [
     "aiocache>=0.12.1,<1.0",
@@ -12,7 +12,7 @@ dependencies = [
     "gunicorn>=21.1.0",
     "huggingface-hub>=0.20.1,<1.0",
     "insightface>=0.7.3,<1.0",
-    "numpy<2",
+    "numpy>=2.3.4",
     "opencv-python-headless>=4.7.0.72,<5.0",
     "orjson>=3.9.5",
     "pillow>=9.5.0,<11.0",
@@ -49,24 +49,16 @@ lint = [
 dev = ["locust>=2.15.1", { include-group = "test" }, { include-group = "lint" }]
 
 [project.optional-dependencies]
-cpu = ["onnxruntime>=1.15.0,<2"]
-cuda = ["onnxruntime-gpu>=1.17.0,<2"]
-openvino = ["onnxruntime-openvino>=1.17.1,<1.19.0"]
-armnn = ["onnxruntime>=1.15.0,<2"]
-rknn = ["onnxruntime>=1.15.0,<2", "rknn-toolkit-lite2>=2.3.0,<3"]
+cpu = ["onnxruntime>=1.23.0,<2"]
+cuda = ["onnxruntime-gpu>=1.23.0,<2"]
+openvino = ["onnxruntime-openvino>=1.23.0,<2"]
+armnn = ["onnxruntime>=1.23.0,<2"]
+rknn = ["onnxruntime>=1.23.0,<2", "rknn-toolkit-lite2>=2.3.0,<3"]
 rocm = []
 
 [tool.uv]
 compile-bytecode = true
 
-[[tool.uv.index]]
-name = "cuda12"
-url = "https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/"
-explicit = true
-
-[tool.uv.sources]
-onnxruntime-gpu = { index = "cuda12" }
-
 [tool.hatch.build.targets.sdist]
 include = ["immich_ml"]
 

machine-learning/test_main.py

@@ -244,13 +244,19 @@ def test_sets_default_provider_options(self, ov_device_ids: list[str]) -> None:
         session = OrtSession(model_path, providers=["OpenVINOExecutionProvider", "CPUExecutionProvider"])
 
         assert session.provider_options == [
-            {"device_type": "GPU.0", "precision": "FP32", "cache_dir": "/cache/ViT-B-32__openai/openvino"},
+            {
+                "device_type": "GPU.0",
+                "precision": "FP32",
+                "cache_dir": "/cache/ViT-B-32__openai/openvino",
+                "load_config": "{\"GPU.0\":{\"CPU_RUNTIME_CACHE_CAPACITY\":\"20\"}}",
+            },
             {"arena_extend_strategy": "kSameAsRequested"},
         ]
 
-    def test_sets_provider_options_for_openvino(self) -> None:
+    def test_sets_provider_options_for_openvino(self, mocker: MockerFixture) -> None:
         model_path = "/cache/ViT-B-32__openai/textual/model.onnx"
         os.environ["MACHINE_LEARNING_DEVICE_ID"] = "1"
+        mocker.patch.object(settings, "openvino_cache_capacity", 10)
 
         session = OrtSession(model_path, providers=["OpenVINOExecutionProvider"])
 
@@ -259,6 +265,7 @@ def test_sets_provider_options_for_openvino(self) -> None:
                 "device_type": "GPU.1",
                 "precision": "FP32",
                 "cache_dir": "/cache/ViT-B-32__openai/textual/openvino",
+                "load_config": "{\"GPU.1\":{\"CPU_RUNTIME_CACHE_CAPACITY\":\"10\"}}"
             }
         ]