From b163ae6a4dcb9b5024888c33620db8ddcd01b037 Mon Sep 17 00:00:00 2001 From: Kevin Turner <83819+keturn@users.noreply.github.com> Date: Fri, 28 Jul 2023 21:30:20 -0700 Subject: [PATCH] refactor(ModelManager): factor out get_model_config --- invokeai/backend/model_management/model_manager.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/invokeai/backend/model_management/model_manager.py b/invokeai/backend/model_management/model_manager.py index 5da4344b89..51053f92cc 100644 --- a/invokeai/backend/model_management/model_manager.py +++ b/invokeai/backend/model_management/model_manager.py @@ -460,7 +460,7 @@ class ModelManager(object): if not self.model_exists(model_name, base_model, model_type, rescan=True): raise ModelNotFoundException(f"Model not found - {model_key}") - model_config = self.models[model_key] + model_config = self._get_model_config(base_model, model_name, model_type) model_path = self.app_config.root_path / model_config.path if not model_path.exists(): @@ -518,6 +518,14 @@ class ModelManager(object): _cache=self.cache, ) + def _get_model_config(self, base_model, model_name, model_type) -> ModelConfigBase: + model_key = self.create_key(model_name, base_model, model_type) + try: + model_config = self.models[model_key] + except KeyError: + raise ModelNotFoundException(f"Model not found - {model_key}") + return model_config + def _get_implementation(self, base_model: BaseModelType, model_type: ModelType) -> type[ModelBase]: model_class = MODEL_CLASSES[base_model][model_type] return model_class