| from queue import Queue |
| from typing import TYPE_CHECKING, Optional, TypeVar |
|
|
| from invokeai.app.services.object_serializer.object_serializer_base import ObjectSerializerBase |
|
|
| T = TypeVar("T") |
|
|
| if TYPE_CHECKING: |
| from invokeai.app.services.invoker import Invoker |
|
|
|
|
| class ObjectSerializerForwardCache(ObjectSerializerBase[T]): |
| """ |
| Provides a LRU cache for an instance of `ObjectSerializerBase`. |
| Saving an object to the cache always writes through to the underlying storage. |
| """ |
|
|
| def __init__(self, underlying_storage: ObjectSerializerBase[T], max_cache_size: int = 20): |
| super().__init__() |
| self._underlying_storage = underlying_storage |
| self._cache: dict[str, T] = {} |
| self._cache_ids = Queue[str]() |
| self._max_cache_size = max_cache_size |
|
|
| def start(self, invoker: "Invoker") -> None: |
| self._invoker = invoker |
| start_op = getattr(self._underlying_storage, "start", None) |
| if callable(start_op): |
| start_op(invoker) |
|
|
| def stop(self, invoker: "Invoker") -> None: |
| self._invoker = invoker |
| stop_op = getattr(self._underlying_storage, "stop", None) |
| if callable(stop_op): |
| stop_op(invoker) |
|
|
| def load(self, name: str) -> T: |
| cache_item = self._get_cache(name) |
| if cache_item is not None: |
| return cache_item |
|
|
| obj = self._underlying_storage.load(name) |
| self._set_cache(name, obj) |
| return obj |
|
|
| def save(self, obj: T) -> str: |
| name = self._underlying_storage.save(obj) |
| self._set_cache(name, obj) |
| return name |
|
|
| def delete(self, name: str) -> None: |
| self._underlying_storage.delete(name) |
| if name in self._cache: |
| del self._cache[name] |
| self._on_deleted(name) |
|
|
| def _get_cache(self, name: str) -> Optional[T]: |
| return None if name not in self._cache else self._cache[name] |
|
|
| def _set_cache(self, name: str, data: T): |
| if name not in self._cache: |
| self._cache[name] = data |
| self._cache_ids.put(name) |
| if self._cache_ids.qsize() > self._max_cache_size: |
| self._cache.pop(self._cache_ids.get()) |
|
|
