MertandGitHub
23f1f87dd0
* add immich_model * graph per shape * decouple ocr from ort * prepared graphs * use f16 for accelerators * preload all shapes * warm-up pass after preload * keep precision for fp32 graphs * ensure padding for text recognition * optimize postprocessing for fp16 outputs * handle raw ctc logits * use f16 output for rknn * mmap token embedding * read device metadata
44 lines
1.5 KiB
Python
44 lines
1.5 KiB
Python
import asyncio
|
|
from typing import Any
|
|
|
|
from immich_ml import allocator
|
|
from immich_ml.models import get_model_class
|
|
from immich_ml.models.base import InferenceModel
|
|
|
|
from ..schemas import ModelTask, ModelType
|
|
|
|
|
|
class ModelCache:
|
|
def __init__(self) -> None:
|
|
self._models: dict[tuple[Any, ...], InferenceModel] = {}
|
|
self._expiries: dict[tuple[Any, ...], asyncio.TimerHandle] = {}
|
|
|
|
def get(
|
|
self, model_name: str, model_type: ModelType, model_task: ModelTask, ttl: int | None = None, **options: Any
|
|
) -> InferenceModel:
|
|
model_cls = get_model_class(model_name, model_type, model_task)
|
|
key = (model_name, model_type, model_task, *(options.get(option) for option in model_cls.graph_options))
|
|
if key not in self._models:
|
|
self._models[key] = model_cls(model_name, **options)
|
|
elif key not in self._expiries:
|
|
ttl = None # a model that was preloaded stays
|
|
if ttl:
|
|
if key in self._expiries:
|
|
self._expiries[key].cancel()
|
|
self._expiries[key] = asyncio.get_running_loop().call_later(ttl, self._evict, key)
|
|
return self._models[key]
|
|
|
|
def clear(self) -> None:
|
|
if not self._models:
|
|
return
|
|
self._models.clear()
|
|
for key in list(self._expiries.keys()):
|
|
self._expiries.pop(key).cancel()
|
|
allocator.release()
|
|
|
|
def _evict(self, key: tuple[Any, ...]) -> None:
|
|
if key not in self._models:
|
|
return
|
|
del self._models[key], self._expiries[key]
|
|
allocator.release()
|