207 lines
7.1 KiB
Python
207 lines
7.1 KiB
Python
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})"
|