helpy4/immich
1
0
Fork 0
mirror of https://github.com/immich-app/immich synced 2025-10-17 18:19:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(ml): coreml (#17718)

Browse source 
* coreml

* add test

* use arena by default in native installation

* fix tests

* add env to docs

* remove availability envs
This commit is contained in:
Mert 2025-10-14 13:51:31 -04:00 • committed by GitHub
parent 43eccca86a
commit 1b62c2ef55
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 47 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

28
machine-learning/test_main.py
View file

@ -180,6 +180,7 @@ class TestOrtSession:
CUDA_EP_OUT_OF_ORDER = ["CPUExecutionProvider", "CUDAExecutionProvider"]
TRT_EP = ["TensorrtExecutionProvider", "CUDAExecutionProvider", "CPUExecutionProvider"]
ROCM_EP = ["ROCMExecutionProvider", "CPUExecutionProvider"]
COREML_EP = ["CoreMLExecutionProvider", "CPUExecutionProvider"]
@pytest.mark.providers(CPU_EP)
def test_sets_cpu_provider(self, providers: list[str]) -> None:
@ -225,6 +226,12 @@ class TestOrtSession:
assert session.providers == self.ROCM_EP
@pytest.mark.providers(COREML_EP)
def test_uses_coreml(self, providers: list[str]) -> None:
session = OrtSession("ViT-B-32__openai")
assert session.providers == self.COREML_EP
def test_sets_provider_kwarg(self) -> None:
providers = ["CUDAExecutionProvider"]
session = OrtSession("ViT-B-32__openai", providers=providers)
@ -284,7 +291,6 @@ class TestOrtSession:
assert session.sess_options.execution_mode == ort.ExecutionMode.ORT_SEQUENTIAL
assert session.sess_options.inter_op_num_threads == 1
assert session.sess_options.intra_op_num_threads == 2
assert session.sess_options.enable_cpu_mem_arena is False
def test_sets_default_sess_options_does_not_set_threads_if_non_cpu_and_default_threads(self) -> None:
session = OrtSession("ViT-B-32__openai", providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
@ -302,6 +308,26 @@ class TestOrtSession:
assert session.sess_options.inter_op_num_threads == 2
assert session.sess_options.intra_op_num_threads == 4
def test_uses_arena_if_enabled(self, mocker: MockerFixture) -> None:
mock_settings = mocker.patch("immich_ml.sessions.ort.settings", autospec=True)
mock_settings.model_inter_op_threads = 0
mock_settings.model_intra_op_threads = 0
mock_settings.model_arena = True
session = OrtSession("ViT-B-32__openai", providers=["CPUExecutionProvider"])
assert session.sess_options.enable_cpu_mem_arena
def test_does_not_use_arena_if_disabled(self, mocker: MockerFixture) -> None:
mock_settings = mocker.patch("immich_ml.sessions.ort.settings", autospec=True)
mock_settings.model_inter_op_threads = 0
mock_settings.model_intra_op_threads = 0
mock_settings.model_arena = False
session = OrtSession("ViT-B-32__openai", providers=["CPUExecutionProvider"])
assert not session.sess_options.enable_cpu_mem_arena
def test_sets_sess_options_kwarg(self) -> None:
sess_options = ort.SessionOptions()
session = OrtSession(