From 97a7f51721548fe2b59d006dd0899da1596add11 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Thu, 18 Jul 2024 15:34:01 -0400 Subject: [PATCH] don't use cpu state_dict for model unpatching when executing on cpu (#6631) Co-authored-by: Lincoln Stein --- .../model_manager/load/model_cache/model_cache_default.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache_default.py b/invokeai/backend/model_manager/load/model_cache/model_cache_default.py index e69201e739..482585e8e7 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache_default.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache_default.py @@ -167,7 +167,8 @@ class ModelCache(ModelCacheBase[AnyModel]): size = calc_model_size_by_data(self.logger, model) self.make_room(size) - state_dict = model.state_dict() if isinstance(model, torch.nn.Module) else None + running_on_cpu = self.execution_device == torch.device("cpu") + state_dict = model.state_dict() if isinstance(model, torch.nn.Module) and not running_on_cpu else None cache_record = CacheRecord(key=key, model=model, device=self.storage_device, state_dict=state_dict, size=size) self._cached_models[key] = cache_record self._cache_stack.append(key)