when model is forcibly moved into RAM update loaded_models set

2024-08-30 20:32:17 +00:00 · 2023-05-04 23:28:03 -04:00 · 2023-05-04 23:28:03 -04:00 · a4e36bc02a
commit a4e36bc02a
parent 68bc0112fa
1 changed files with 1 additions and 0 deletions
--- a/invokeai/backend/model_management/model_cache.py
+++ b/invokeai/backend/model_management/model_cache.py
@ -178,6 +178,7 @@ class ModelCache(object):
            if hasattr(model,'to') and (key in self.loaded_models
                                        and self.locked_models[key] == 0):
                model.to(self.storage_device)
+                self.loaded_models.remove(key)
            yield model

    def _offload_unlocked_models(self):