from __future__ import annotations from abc import ABC, abstractmethod from collections.abc import Hashable from dataclasses import dataclass from pathlib import Path from shutil import rmtree from threading import Lock from typing import Any, ClassVar, Self from huggingface_hub import snapshot_download import immich_ml.sessions.ann.loader import immich_ml.sessions.rknn as rknn from immich_ml.sessions.ort import OrtSession from ..config import clean_name, log, settings from ..schemas import ModelFormat, ModelIdentity, ModelSession, ModelSource, ModelTask, ModelType, Options from ..sessions.ann import AnnSession from ..sessions.policy import ShapePolicy from .constants import get_model_source _IGNORED_PATTERNS: dict[ModelFormat, list[str]] = { ModelFormat.ONNX: ["*.armnn", "*.rknn"], ModelFormat.ARMNN: ["*.rknn"], ModelFormat.RKNN: ["*.armnn"], } class InferenceModel[O: Options](ABC): depends: ClassVar[list[ModelIdentity]] identity: ClassVar[ModelIdentity] sources: ClassVar[tuple[ModelSource, ...]] # the shapes this model feeds, which pick the artifact and tell the engine what it may pin shape_policy: ShapePolicy = ShapePolicy() threads: ClassVar[int] = 2 # how many threads one run takes for CPU inference def __init__( self, model_name: str, cache_dir: Path | str | None = None, model_format: ModelFormat | None = None, session: ModelSession | None = None, ) -> None: self.loaded = session is not None self.load_attempts = 0 self._load_lock = Lock() 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 = model_format if model_format is not None else self._model_format_default if session is not None: self.session = session @classmethod def graph(cls, options: O) -> Hashable: """What of its options the graph is built for: an instance serves only requests that agree on it.""" return None @classmethod def create(cls, model_name: str, options: O) -> Self: return cls(model_name) def download(self) -> None: if self.cached: return model_type = self.model_type.replace("-", " ") log.info(f"Downloading {model_type} model '{self.model_name}' to {self.model_dir}. This may take a while.") snapshot_download( f"{settings.model_organization}/{self.model_name}", revision=settings.model_revision, cache_dir=self.cache_dir, local_dir=self.cache_dir, ignore_patterns=_IGNORED_PATTERNS.get(self.model_format, []), ) if not self.cached: # the repository has nothing in this format, which is what sends a model back to ONNX raise FileNotFoundError(f"Model file not found: {self.model_path}") def load(self) -> None: with self._load_lock: # two requests may arrive for a model neither has loaded if self.loaded: return self.load_attempts += 1 self.download() 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 def unload(self) -> None: with self._load_lock: if self.loaded: del self.session self.loaded = False def build(self) -> None: self.load() self.session.warm() def predict(self, *inputs: Any, options: O) -> Any: self.load() return self._predict(*inputs, options=options) @abstractmethod def _predict(self, *inputs: Any, **model_kwargs: Any) -> Any: ... def _load(self) -> ModelSession: return self._make_session() def clear_cache(self) -> None: if not self.cache_dir.exists(): log.warning( f"Attempted to clear cache for model '{self.model_name}', but cache directory does not exist", ) return if not rmtree.avoids_symlink_attacks: raise RuntimeError("Attempted to clear cache, but rmtree is not safe on this platform") if self.cache_dir.is_dir(): log.info(f"Cleared cache directory for model '{self.model_name}'.") rmtree(self.cache_dir) else: log.warning( ( f"Encountered file instead of directory at cache path " f"for '{self.model_name}'. Removing file and replacing with a directory." ), ) self.cache_dir.unlink() self.cache_dir.mkdir(parents=True, exist_ok=True) def _make_session(self) -> ModelSession: match self.model_format: case ModelFormat.ARMNN: return AnnSession(self.model_path) case ModelFormat.ONNX: return OrtSession(self.model_path, self.shape_policy, threads=self.threads) case ModelFormat.RKNN: return rknn.RknnSession(self.model_path) @property def model_dir(self) -> Path: return self.cache_dir / self.model_type.value @property def model_path(self) -> Path: if self.model_format == ModelFormat.RKNN: return rknn.model_path(self.model_dir, self.shape_policy.label) return self.model_dir / f"model.{self.model_format}" @property def model_task(self) -> ModelTask: return self.identity[1] @property def model_type(self) -> ModelType: return self.identity[0] @property def cache_dir(self) -> Path: return self._cache_dir @cache_dir.setter def cache_dir(self, cache_dir: Path) -> None: self._cache_dir = cache_dir @property def _cache_dir_default(self) -> Path: cache_dir = settings.cache_folder / self.model_task.value / self.model_name return cache_dir if settings.legacy_models else cache_dir / settings.model_revision @property def cached(self) -> bool: return self.model_path.is_file() @property def model_format(self) -> ModelFormat: return self._model_format @model_format.setter 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: if rknn.is_available: return ModelFormat.RKNN elif immich_ml.sessions.ann.loader.is_available and settings.ann: return ModelFormat.ARMNN else: return ModelFormat.ONNX @dataclass(frozen=True) class InferenceEntry[O: Options]: model: type[InferenceModel[O]] name: str options: O def __post_init__(self) -> None: if get_model_source(self.name) not in self.model.sources: raise ValueError(f"'{self.name}' is not a model {self.model.__name__} runs") def __str__(self) -> str: return f"{self.model.__name__}(name={self.name!r}, options={self.options})"