feat(ml): add preload and fp16 settings for ocr (#23576)

machine-learning/immich_ml/config.py

@@ -13,6 +13,8 @@ from rich.logging import RichHandler
 from uvicorn import Server
 from uvicorn.workers import UvicornWorker
 
+from .schemas import ModelPrecision
+
 
 class ClipSettings(BaseModel):
     textual: str | None = None
@@ -24,6 +26,11 @@ class FacialRecognitionSettings(BaseModel):
     detection: str | None = None
 
 
+class OcrSettings(BaseModel):
+    recognition: str | None = None
+    detection: str | None = None
+
+
 class PreloadModelData(BaseModel):
     clip_fallback: str | None = os.getenv("MACHINE_LEARNING_PRELOAD__CLIP", None)
     facial_recognition_fallback: str | None = os.getenv("MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION", None)
@@ -37,6 +44,7 @@ class PreloadModelData(BaseModel):
         del os.environ["MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION"]
     clip: ClipSettings = ClipSettings()
     facial_recognition: FacialRecognitionSettings = FacialRecognitionSettings()
+    ocr: OcrSettings = OcrSettings()
 
 
 class MaxBatchSize(BaseModel):
@@ -70,6 +78,7 @@ class Settings(BaseSettings):
     rknn_threads: int = 1
     preload: PreloadModelData | None = None
     max_batch_size: MaxBatchSize | None = None
+    openvino_precision: ModelPrecision = ModelPrecision.FP32
 
     @property
     def device_id(self) -> str:

machine-learning/immich_ml/main.py

@@ -103,6 +103,20 @@ async def preload_models(preload: PreloadModelData) -> None:
             ModelTask.FACIAL_RECOGNITION,
         )
 
+    if preload.ocr.detection is not None:
+        await load_models(
+            preload.ocr.detection,
+            ModelType.DETECTION,
+            ModelTask.OCR,
+        )
+
+    if preload.ocr.recognition is not None:
+        await load_models(
+            preload.ocr.recognition,
+            ModelType.RECOGNITION,
+            ModelTask.OCR,
+        )
+
     if preload.clip_fallback is not None:
         log.warning(
             "Deprecated env variable: 'MACHINE_LEARNING_PRELOAD__CLIP'. "

machine-learning/immich_ml/schemas.py

@@ -46,6 +46,11 @@ class ModelSource(StrEnum):
     PADDLE = "paddle"
 
 
+class ModelPrecision(StrEnum):
+    FP16 = "FP16"
+    FP32 = "FP32"
+
+
 ModelIdentity = tuple[ModelType, ModelTask]
 
 

machine-learning/immich_ml/sessions/ort.py

@@ -93,10 +93,12 @@ class OrtSession:
                 case "CUDAExecutionProvider" | "ROCMExecutionProvider":
                     options = {"arena_extend_strategy": "kSameAsRequested", "device_id": settings.device_id}
                 case "OpenVINOExecutionProvider":
+                    openvino_dir = self.model_path.parent / "openvino"
+                    device = f"GPU.{settings.device_id}"
                     options = {
-                        "device_type": f"GPU.{settings.device_id}",
-                        "precision": "FP32",
-                        "cache_dir": (self.model_path.parent / "openvino").as_posix(),
+                        "device_type": device,
+                        "precision": settings.openvino_precision.value,
+                        "cache_dir": openvino_dir.as_posix(),
                     }
                 case "CoreMLExecutionProvider":
                     options = {