From 06510bd75572998818b4428b605dcec0db106e5e Mon Sep 17 00:00:00 2001
From: mertalev <101130780+mertalev@users.noreply.github.com>
Date: Thu, 19 Dec 2024 20:53:40 -0500
Subject: [PATCH] bump versions, run on mich

---
 .github/workflows/docker.yml             | 4 ++--
 machine-learning/Dockerfile              | 8 ++++----
 machine-learning/app/models/constants.py | 7 ++++++-
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml
index a4b1a0dbd4..4f3e7185ce 100644
--- a/.github/workflows/docker.yml
+++ b/.github/workflows/docker.yml
@@ -48,7 +48,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        suffix: ['', '-cuda', '-openvino', '-armnn']
+        suffix: ['', '-cuda', '-rocm', '-openvino', '-armnn']
     steps:
       - name: Login to GitHub Container Registry
         uses: docker/login-action@v3
@@ -91,7 +91,7 @@ jobs:
     name: Build and Push ML
     needs: pre-job
     if: ${{ needs.pre-job.outputs.should_run_ml == 'true' }}
-    runs-on: ubuntu-latest
+    runs-on: mich
     env:
       image: immich-machine-learning
       context: machine-learning
diff --git a/machine-learning/Dockerfile b/machine-learning/Dockerfile
index 7c852e9622..e3287edc7d 100644
--- a/machine-learning/Dockerfile
+++ b/machine-learning/Dockerfile
@@ -17,7 +17,7 @@ RUN mkdir /opt/armnn && \
 
 # Warning: 26.3Gb of disk space required to pull this image
 # https://github.com/microsoft/onnxruntime/blob/main/dockerfiles/Dockerfile.rocm
-FROM rocm/dev-ubuntu-22.04:6.1.2-complete as builder-rocm
+FROM rocm/dev-ubuntu-24.04:6.2.4-complete AS builder-rocm
 
 WORKDIR /code
 
@@ -32,7 +32,7 @@ RUN wget https://github.com/Kitware/CMake/releases/download/v3.27.3/cmake-3.27.3
 ENV PATH /code/cmake-3.27.3-linux-x86_64/bin:${PATH}
 
 # Prepare onnxruntime repository & build onnxruntime
-RUN git clone --single-branch --branch v1.18.1 --recursive "https://github.com/Microsoft/onnxruntime" onnxruntime
+RUN git clone --single-branch --branch v1.20.1 --recursive "https://github.com/Microsoft/onnxruntime" onnxruntime
 WORKDIR /code/onnxruntime
 # EDIT PR
 # While there's still this PR open, we need to compile on the branch of the PR
@@ -46,7 +46,7 @@ RUN /bin/sh ./dockerfiles/scripts/install_common_deps.sh
 # I lowered the number of threads to 8, and it worked
 # Even with 12 threads, the compilation took more than 1,5 hours to fail
 RUN ./build.sh --allow_running_as_root --config Release --build_wheel --update --build --parallel 9 --cmake_extra_defines\
-    ONNXRUNTIME_VERSION=1.18.1 --use_rocm --rocm_home=/opt/rocm
+    ONNXRUNTIME_VERSION=1.20.1 --use_rocm --rocm_home=/opt/rocm
 RUN mv /code/onnxruntime/build/Linux/Release/dist/*.whl /opt/
 
 FROM builder-${DEVICE} AS builder
@@ -117,7 +117,7 @@ COPY --from=builder-armnn \
     /opt/ann/build.sh \
     /opt/armnn/
 
-FROM rocm/dev-ubuntu-22.04:6.1.2-complete AS prod-rocm
+FROM rocm/dev-ubuntu-24.04:6.2.4-complete AS prod-rocm
 
 
 FROM prod-${DEVICE} AS prod
diff --git a/machine-learning/app/models/constants.py b/machine-learning/app/models/constants.py
index d5b3f4dfd2..43088741b9 100644
--- a/machine-learning/app/models/constants.py
+++ b/machine-learning/app/models/constants.py
@@ -63,7 +63,12 @@ _INSIGHTFACE_MODELS = {
 }
 
 
-SUPPORTED_PROVIDERS = ["CUDAExecutionProvider", "ROCMExecutionProvider", "OpenVINOExecutionProvider", "CPUExecutionProvider"]
+SUPPORTED_PROVIDERS = [
+    "CUDAExecutionProvider",
+    "ROCMExecutionProvider",
+    "OpenVINOExecutionProvider",
+    "CPUExecutionProvider",
+]
 
 
 def get_model_source(model_name: str) -> ModelSource | None: