fix(ml): armnn not being used (#10929)

* fix armnn not being used, move fallback handling to main, add tests

* formatting

machine-learning/app/models/base.py

@@ -23,7 +23,7 @@ class InferenceModel(ABC):
         self,
         model_name: str,
         cache_dir: Path | str | None = None,
-        preferred_format: ModelFormat | None = None,
+        model_format: ModelFormat | None = None,
         session: ModelSession | None = None,
         **model_kwargs: Any,
     ) -> None:
@@ -31,7 +31,7 @@ class InferenceModel(ABC):
         self.load_attempts = 0
         self.model_name = clean_name(model_name)
         self.cache_dir = Path(cache_dir) if cache_dir is not None else self._cache_dir_default
-        self.model_format = preferred_format if preferred_format is not None else self._model_format_default
+        self.model_format = model_format if model_format is not None else self._model_format_default
         if session is not None:
             self.session = session
 
@@ -48,7 +48,7 @@ class InferenceModel(ABC):
         self.load_attempts += 1
 
         self.download()
-        attempt = f"Attempt #{self.load_attempts + 1} to load" if self.load_attempts else "Loading"
+        attempt = f"Attempt #{self.load_attempts} to load" if self.load_attempts > 1 else "Loading"
         log.info(f"{attempt} {self.model_type.replace('-', ' ')} model '{self.model_name}' to memory")
         self.session = self._load()
         self.loaded = True
@@ -101,6 +101,9 @@ class InferenceModel(ABC):
         self.cache_dir.mkdir(parents=True, exist_ok=True)
 
     def _make_session(self, model_path: Path) -> ModelSession:
+        if not model_path.is_file():
+            raise FileNotFoundError(f"Model file not found: {model_path}")
+
         match model_path.suffix:
             case ".armnn":
                 session: ModelSession = AnnSession(model_path)
@@ -144,17 +147,13 @@ class InferenceModel(ABC):
 
     @property
     def model_format(self) -> ModelFormat:
-        return self._preferred_format
+        return self._model_format
 
     @model_format.setter
-    def model_format(self, preferred_format: ModelFormat) -> None:
-        log.debug(f"Setting preferred format to {preferred_format}")
-        self._preferred_format = preferred_format
+    def model_format(self, model_format: ModelFormat) -> None:
+        log.debug(f"Setting model format to {model_format}")
+        self._model_format = model_format
 
     @property
     def _model_format_default(self) -> ModelFormat:
-        prefer_ann = ann.ann.is_available and settings.ann
-        ann_exists = (self.model_dir / "model.armnn").is_file()
-        if prefer_ann and not ann_exists:
-            log.warning(f"ARM NN is available, but '{self.model_name}' does not support ARM NN. Falling back to ONNX.")
-        return ModelFormat.ARMNN if prefer_ann and ann_exists else ModelFormat.ONNX
+        return ModelFormat.ARMNN if ann.ann.is_available and settings.ann else ModelFormat.ONNX

machine-learning/app/models/clip/textual.py

@@ -22,11 +22,12 @@ class BaseCLIPTextualEncoder(InferenceModel):
         return res
 
     def _load(self) -> ModelSession:
+        session = super()._load()
         log.debug(f"Loading tokenizer for CLIP model '{self.model_name}'")
         self.tokenizer = self._load_tokenizer()
         log.debug(f"Loaded tokenizer for CLIP model '{self.model_name}'")
 
-        return super()._load()
+        return session
 
     @abstractmethod
     def _load_tokenizer(self) -> Tokenizer:

machine-learning/app/models/facial_recognition/detection.py

@@ -1,4 +1,3 @@
-from pathlib import Path
 from typing import Any
 
 import numpy as np
@@ -14,15 +13,9 @@ class FaceDetector(InferenceModel):
     depends = []
     identity = (ModelType.DETECTION, ModelTask.FACIAL_RECOGNITION)
 
-    def __init__(
-        self,
-        model_name: str,
-        min_score: float = 0.7,
-        cache_dir: Path | str | None = None,
-        **model_kwargs: Any,
-    ) -> None:
+    def __init__(self, model_name: str, min_score: float = 0.7, **model_kwargs: Any) -> None:
         self.min_score = model_kwargs.pop("minScore", min_score)
-        super().__init__(model_name, cache_dir, **model_kwargs)
+        super().__init__(model_name, **model_kwargs)
 
     def _load(self) -> ModelSession:
         session = self._make_session(self.model_path)

machine-learning/app/models/facial_recognition/recognition.py

@@ -9,7 +9,7 @@ from numpy.typing import NDArray
 from onnx.tools.update_model_dims import update_inputs_outputs_dims
 from PIL import Image
 
-from app.config import clean_name, log
+from app.config import log
 from app.models.base import InferenceModel
 from app.models.transforms import decode_cv2
 from app.schemas import FaceDetectionOutput, FacialRecognitionOutput, ModelFormat, ModelSession, ModelTask, ModelType
@@ -20,20 +20,14 @@ class FaceRecognizer(InferenceModel):
     depends = [(ModelType.DETECTION, ModelTask.FACIAL_RECOGNITION)]
     identity = (ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION)
 
-    def __init__(
-        self,
-        model_name: str,
-        min_score: float = 0.7,
-        cache_dir: Path | str | None = None,
-        **model_kwargs: Any,
-    ) -> None:
-        super().__init__(clean_name(model_name), cache_dir, **model_kwargs)
+    def __init__(self, model_name: str, min_score: float = 0.7, **model_kwargs: Any) -> None:
+        super().__init__(model_name, **model_kwargs)
         self.min_score = model_kwargs.pop("minScore", min_score)
         self.batch = self.model_format == ModelFormat.ONNX
 
     def _load(self) -> ModelSession:
         session = self._make_session(self.model_path)
-        if self.model_format == ModelFormat.ONNX and not has_batch_axis(session):
+        if self.batch and not has_batch_axis(session):
             self._add_batch_axis(self.model_path)
             session = self._make_session(self.model_path)
         self.model = ArcFaceONNX(