[mertalev]

Description

While facial recognition already used ONNX models, this change brings them over to other ML tasks as well. There are a number of important advantages with this change:

1. Higher performance: ONNX Runtime is a faster engine than PyTorch.
2. Single runtime for all models: Since everything will use ONNX Runtime, changes can be applied to all models without any custom logic.
3. Supports many backends: Beyond CPU and CUDA, ORT has support for too many backends to list. This enables future improvements.
4. More customizable: It allows for deep customization of threading and memory management, allowing it to be better-optimized for Immich.
5. Better memory management: Related to (4), ORT makes it easier to keep RAM usage on a leash, and has better behavior when models are unloaded and reloaded.
6. More model options. Users can choose among OpenCLIP and M-CLIP models as well, with the latter allowing for a better quality multilingual model.

NOTE: There's a breaking change in this PR. Since Sentence-Transformers is no longer used for CLIP models, users who set MACHINE_LEARNING_CLIP_IMAGE_MODEL or MACHINE_LEARNING_CLIP_TEXT_MODEL must migrate to one of the ViT-B models listed here, with the caveat that OpenCLIP models will require running CLIP on all images since the embeddings are incompatible.

How Has This Been Tested?

I deleted the model cache volume and started the server to confirm that it downloaded each model on startup without error. After the ML service completed startup, I started jobs for object tagging, CLIP and facial recognition and observed no errors. Performance is higher than on main while RAM consumption is about equal.

Edit: I did some load-testing with Locust for more formal results. ML on this branch is over 3.5x faster compared to main.

Likewise, latency is much lower as well.

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

1 Ignored Deployment

Name	Status	Preview	Comments	Updated (UTC)
immich	⬜️ Ignored (Inspect)	Visit Preview		Aug 21, 2023 2:43am

[alextran1502]

Hello, for people for existing model_cache volume, I assume it will download the new models and use the new model to run, correct?

[mertalev]

Yup, it'll download a model if it doesn't see it in the cache, so the model cache doesn't need any user intervention. But the old models will still be there until they're deleted by the user.

[LeezPayne]

Great job! I'm wondering if this would be better that users can manually configure the expected size of the CLIP model output dimensions listed here, as some may want to use a different one, such as M-CLIP/XLM-Roberta-Large-Vit-L-14 with 768 output dimensions. Here are some codes that may be related.

immich/server/src/infra/typesense-schemas/asset.schema.ts

Line 31 in 165b91b

{ name: 'smartInfo.clipEmbedding', type: 'float[]', facet: false, optional: true, num_dim: 512 },

immich/machine-learning/app/test_main.py

Line 75 in 165b91b

embedding = np.random.rand(512).astype(np.float32)

[mertalev]

Thanks! I'm planning on moving the ML settings to the admin page so they're stored in the server and sent as part of the request. That way, the server can know what embedding size it should set instead of hardcoding it.

[jrasm91]

Although, to some extent that specific variable might require some more work. There are other places where the length is specifically set, like in the typesense schema, it has a hardcoded length of 512

[mertalev]

Yeah, I think I'd have to update the asset schema to use the new embedding size and drop the current embeddings. (And give a pop-up warning about this beforehand).

Hmm maybe it'd be better to do this after switching away from Typesense since I'd have to rewrite this logic for pgvector.

[jrasm91]

Yeah. I think we can get all the rest of the options configurable pretty easily though, so we can definitely start there.

[waclaw66]

Isn't clip_server missing within python dependencies?

srp 26 09:14:32 server immich-machine-learning[440443]: Traceback (most recent call last):
srp 26 09:14:32 server immich-machine-learning[440443]:   File "<frozen runpy>", line 198, in _run_module_as_main
srp 26 09:14:32 server immich-machine-learning[440443]:   File "<frozen runpy>", line 88, in _run_code
srp 26 09:14:32 server immich-machine-learning[440443]:   File "/opt/immich/machine-learning/app/main.py", line 13, in <module>
srp 26 09:14:32 server immich-machine-learning[440443]:     from app.models.base import InferenceModel
srp 26 09:14:32 server immich-machine-learning[440443]:   File "/opt/immich/machine-learning/app/models/__init__.py", line 1, in <module>
srp 26 09:14:32 server immich-machine-learning[440443]:     from .clip import CLIPEncoder
srp 26 09:14:32 server immich-machine-learning[440443]:   File "/opt/immich/machine-learning/app/models/clip.py", line 7, in <module>
srp 26 09:14:32 server immich-machine-learning[440443]:     from clip_server.model.clip import BICUBIC, _convert_image_to_rgb
srp 26 09:14:32 server immich-machine-learning[440443]: ModuleNotFoundError: No module named 'clip_server'

[mertalev]

clip-server is installed separately in the Dockerfile with the --no-deps flag since it has a lot of unnecessary dependencies.