42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
import asyncio
|
|
from collections.abc import Hashable
|
|
from typing import Any
|
|
|
|
from immich_ml import allocator
|
|
from immich_ml.models.base import InferenceEntry, InferenceModel
|
|
from immich_ml.schemas import Options
|
|
|
|
Key = tuple[type[InferenceModel[Any]], str, Hashable]
|
|
|
|
|
|
class ModelCache:
|
|
def __init__(self) -> None:
|
|
self._models: dict[Key, InferenceModel[Any]] = {}
|
|
self._expiries: dict[Key, asyncio.TimerHandle] = {}
|
|
|
|
def get[O: Options](self, entry: InferenceEntry[O], ttl: int | None = None) -> InferenceModel[O]:
|
|
key = (entry.model, entry.name, entry.model.graph(entry.options))
|
|
if key not in self._models:
|
|
self._models[key] = entry.model.create(entry.name, entry.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: Key) -> None:
|
|
if key not in self._models:
|
|
return
|
|
del self._models[key], self._expiries[key]
|
|
allocator.release()
|