From c579a218ef9253909a9c9648de526c39545b05b7 Mon Sep 17 00:00:00 2001 From: Ryan Dick Date: Mon, 6 Jan 2025 16:48:05 +0000 Subject: [PATCH] Allow models to be locked in VRAM, even if they have been dropped from the RAM cache (related: https://github.com/invoke-ai/InvokeAI/issues/7513). --- .../backend/model_manager/load/load_base.py | 8 ++++---- .../load/model_cache/model_cache.py | 20 +++++++++++++++---- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/invokeai/backend/model_manager/load/load_base.py b/invokeai/backend/model_manager/load/load_base.py index ceea7341499..d62db363a6d 100644 --- a/invokeai/backend/model_manager/load/load_base.py +++ b/invokeai/backend/model_manager/load/load_base.py @@ -57,20 +57,20 @@ def __init__(self, cache_record: CacheRecord, cache: ModelCache): self._cache = cache def __enter__(self) -> AnyModel: - self._cache.lock(self._cache_record.key) + self._cache.lock(self._cache_record) return self.model def __exit__(self, *args: Any, **kwargs: Any) -> None: - self._cache.unlock(self._cache_record.key) + self._cache.unlock(self._cache_record) @contextmanager def model_on_device(self) -> Generator[Tuple[Optional[Dict[str, torch.Tensor]], AnyModel], None, None]: """Return a tuple consisting of the model's state dict (if it exists) and the locked model on execution device.""" - self._cache.lock(self._cache_record.key) + self._cache.lock(self._cache_record) try: yield (self._cache_record.state_dict, self._cache_record.model) finally: - self._cache.unlock(self._cache_record.key) + self._cache.unlock(self._cache_record) @property def model(self) -> AnyModel: diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index dbc3670c95c..cd296aa7bd7 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -194,9 +194,15 @@ def get( return cache_entry - def lock(self, key: str) -> None: + def lock(self, cache_entry: CacheRecord) -> None: """Lock a model for use and move it into VRAM.""" - cache_entry = self._cached_models[key] + if cache_entry.key not in self._cached_models: + self._logger.info( + f"Locking model cache entry {cache_entry.key} ({cache_entry.model.__class__.__name__}), but it has " + "already been dropped from the RAM cache. This is a sign that the model loading order is non-optimal " + "in the invocation code (See https://github.com/invoke-ai/InvokeAI/issues/7513)." + ) + # cache_entry = self._cached_models[key] cache_entry.lock() try: @@ -214,9 +220,15 @@ def lock(self, key: str) -> None: cache_entry.unlock() raise - def unlock(self, key: str) -> None: + def unlock(self, cache_entry: CacheRecord) -> None: """Unlock a model.""" - cache_entry = self._cached_models[key] + if cache_entry.key not in self._cached_models: + self._logger.info( + f"Unlocking model cache entry {cache_entry.key} ({cache_entry.model.__class__.__name__}), but it has " + "already been dropped from the RAM cache. This is a sign that the model loading order is non-optimal " + "in the invocation code (See https://github.com/invoke-ai/InvokeAI/issues/7513)." + ) + # cache_entry = self._cached_models[key] cache_entry.unlock() if not self._lazy_offloading: self._offload_unlocked_models(0)