From 5927ab9c368b9c1e859902a84d0404e29d4bae00 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Sun, 11 Feb 2024 09:27:57 +1100 Subject: [PATCH] chore(backend): rename `ModelInfo` -> `LoadedModelInfo` We have two different classes named `ModelInfo` which might need to be used by API consumers. We need to export both but have to deal with this naming collision. The `ModelInfo` I've renamed here is the one that is returned when a model is loaded. It's the object least likely to be used by API consumers. --- invokeai/app/services/events/events_base.py | 10 +++++----- .../services/model_manager/model_manager_base.py | 4 ++-- .../model_manager/model_manager_default.py | 16 ++++++++-------- .../app/services/shared/invocation_context.py | 7 ++++--- invokeai/backend/__init__.py | 9 ++++++++- invokeai/backend/model_management/__init__.py | 2 +- .../backend/model_management/model_manager.py | 6 +++--- invokeai/backend/util/test_utils.py | 10 +++++----- invokeai/invocation_api/__init__.py | 4 ++-- 9 files changed, 38 insertions(+), 30 deletions(-) diff --git a/invokeai/app/services/events/events_base.py b/invokeai/app/services/events/events_base.py index ad08ae0395..6b441efc2b 100644 --- a/invokeai/app/services/events/events_base.py +++ b/invokeai/app/services/events/events_base.py @@ -11,7 +11,7 @@ from invokeai.app.services.session_queue.session_queue_common import ( SessionQueueStatus, ) from invokeai.app.util.misc import get_timestamp -from invokeai.backend.model_management.model_manager import ModelInfo +from invokeai.backend.model_management.model_manager import LoadedModelInfo from invokeai.backend.model_management.models.base import BaseModelType, ModelType, SubModelType @@ -201,7 +201,7 @@ class EventServiceBase: base_model: BaseModelType, model_type: ModelType, submodel: SubModelType, - model_info: ModelInfo, + loaded_model_info: LoadedModelInfo, ) -> None: """Emitted when a model is correctly loaded (returns model info)""" self.__emit_queue_event( @@ -215,9 +215,9 @@ class EventServiceBase: "base_model": base_model, "model_type": model_type, "submodel": submodel, - "hash": model_info.hash, - "location": str(model_info.location), - "precision": str(model_info.precision), + "hash": loaded_model_info.hash, + "location": str(loaded_model_info.location), + "precision": str(loaded_model_info.precision), }, ) diff --git a/invokeai/app/services/model_manager/model_manager_base.py b/invokeai/app/services/model_manager/model_manager_base.py index a9b53ae224..f888c0ec97 100644 --- a/invokeai/app/services/model_manager/model_manager_base.py +++ b/invokeai/app/services/model_manager/model_manager_base.py @@ -14,8 +14,8 @@ from invokeai.app.services.shared.invocation_context import InvocationContextDat from invokeai.backend.model_management import ( AddModelResult, BaseModelType, + LoadedModelInfo, MergeInterpolationMethod, - ModelInfo, ModelType, SchedulerPredictionType, SubModelType, @@ -48,7 +48,7 @@ class ModelManagerServiceBase(ABC): model_type: ModelType, submodel: Optional[SubModelType] = None, context_data: Optional[InvocationContextData] = None, - ) -> ModelInfo: + ) -> LoadedModelInfo: """Retrieve the indicated model with name and type. submodel can be used to get a part (such as the vae) of a diffusers pipeline.""" diff --git a/invokeai/app/services/model_manager/model_manager_default.py b/invokeai/app/services/model_manager/model_manager_default.py index b641dd3f1e..c3712abf8e 100644 --- a/invokeai/app/services/model_manager/model_manager_default.py +++ b/invokeai/app/services/model_manager/model_manager_default.py @@ -16,8 +16,8 @@ from invokeai.app.services.shared.invocation_context import InvocationContextDat from invokeai.backend.model_management import ( AddModelResult, BaseModelType, + LoadedModelInfo, MergeInterpolationMethod, - ModelInfo, ModelManager, ModelMerger, ModelNotFoundException, @@ -98,7 +98,7 @@ class ModelManagerService(ModelManagerServiceBase): model_type: ModelType, submodel: Optional[SubModelType] = None, context_data: Optional[InvocationContextData] = None, - ) -> ModelInfo: + ) -> LoadedModelInfo: """ Retrieve the indicated model. submodel can be used to get a part (such as the vae) of a diffusers mode. @@ -114,7 +114,7 @@ class ModelManagerService(ModelManagerServiceBase): submodel=submodel, ) - model_info = self.mgr.get_model( + loaded_model_info = self.mgr.get_model( model_name, base_model, model_type, @@ -128,10 +128,10 @@ class ModelManagerService(ModelManagerServiceBase): base_model=base_model, model_type=model_type, submodel=submodel, - model_info=model_info, + loaded_model_info=loaded_model_info, ) - return model_info + return loaded_model_info def model_exists( self, @@ -273,7 +273,7 @@ class ModelManagerService(ModelManagerServiceBase): base_model: BaseModelType, model_type: ModelType, submodel: Optional[SubModelType] = None, - model_info: Optional[ModelInfo] = None, + loaded_model_info: Optional[LoadedModelInfo] = None, ): if self._invoker is None: return @@ -281,7 +281,7 @@ class ModelManagerService(ModelManagerServiceBase): if self._invoker.services.queue.is_canceled(context_data.session_id): raise CanceledException() - if model_info: + if loaded_model_info: self._invoker.services.events.emit_model_load_completed( queue_id=context_data.queue_id, queue_item_id=context_data.queue_item_id, @@ -291,7 +291,7 @@ class ModelManagerService(ModelManagerServiceBase): base_model=base_model, model_type=model_type, submodel=submodel, - model_info=model_info, + loaded_model_info=loaded_model_info, ) else: self._invoker.services.events.emit_model_load_started( diff --git a/invokeai/app/services/shared/invocation_context.py b/invokeai/app/services/shared/invocation_context.py index 1ca44b7862..68fb78c143 100644 --- a/invokeai/app/services/shared/invocation_context.py +++ b/invokeai/app/services/shared/invocation_context.py @@ -13,7 +13,7 @@ from invokeai.app.services.images.images_common import ImageDTO from invokeai.app.services.invocation_services import InvocationServices from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID from invokeai.app.util.step_callback import stable_diffusion_step_callback -from invokeai.backend.model_management.model_manager import ModelInfo +from invokeai.backend.model_management.model_manager import LoadedModelInfo from invokeai.backend.model_management.models.base import BaseModelType, ModelType, SubModelType from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData @@ -272,14 +272,15 @@ class ModelsInterface(InvocationContextInterface): def load( self, model_name: str, base_model: BaseModelType, model_type: ModelType, submodel: Optional[SubModelType] = None - ) -> ModelInfo: + ) -> LoadedModelInfo: """ - Loads a model, returning its `ModelInfo` object. + Loads a model. :param model_name: The name of the model to get. :param base_model: The base model of the model to get. :param model_type: The type of the model to get. :param submodel: The submodel of the model to get. + :returns: An object representing the loaded model. """ # The model manager emits events as it loads the model. It needs the context data to build diff --git a/invokeai/backend/__init__.py b/invokeai/backend/__init__.py index ae9a12edbe..54a1843d46 100644 --- a/invokeai/backend/__init__.py +++ b/invokeai/backend/__init__.py @@ -1,5 +1,12 @@ """ Initialization file for invokeai.backend """ -from .model_management import BaseModelType, ModelCache, ModelInfo, ModelManager, ModelType, SubModelType # noqa: F401 +from .model_management import ( # noqa: F401 + BaseModelType, + LoadedModelInfo, + ModelCache, + ModelManager, + ModelType, + SubModelType, +) from .model_management.models import SilenceWarnings # noqa: F401 diff --git a/invokeai/backend/model_management/__init__.py b/invokeai/backend/model_management/__init__.py index 03abf58eb4..d523a7a0c8 100644 --- a/invokeai/backend/model_management/__init__.py +++ b/invokeai/backend/model_management/__init__.py @@ -3,7 +3,7 @@ Initialization file for invokeai.backend.model_management """ # This import must be first -from .model_manager import AddModelResult, ModelInfo, ModelManager, SchedulerPredictionType +from .model_manager import AddModelResult, LoadedModelInfo, ModelManager, SchedulerPredictionType from .lora import ModelPatcher, ONNXModelPatcher from .model_cache import ModelCache diff --git a/invokeai/backend/model_management/model_manager.py b/invokeai/backend/model_management/model_manager.py index 362d8d3ff5..da74ca3fb5 100644 --- a/invokeai/backend/model_management/model_manager.py +++ b/invokeai/backend/model_management/model_manager.py @@ -271,7 +271,7 @@ CONFIG_FILE_VERSION = "3.0.0" @dataclass -class ModelInfo: +class LoadedModelInfo: context: ModelLocker name: str base_model: BaseModelType @@ -450,7 +450,7 @@ class ModelManager(object): base_model: BaseModelType, model_type: ModelType, submodel_type: Optional[SubModelType] = None, - ) -> ModelInfo: + ) -> LoadedModelInfo: """Given a model named identified in models.yaml, return an ModelInfo object describing it. :param model_name: symbolic name of the model in models.yaml @@ -508,7 +508,7 @@ class ModelManager(object): model_hash = "" # TODO: - return ModelInfo( + return LoadedModelInfo( context=model_context, name=model_name, base_model=base_model, diff --git a/invokeai/backend/util/test_utils.py b/invokeai/backend/util/test_utils.py index 09b9de9e98..685603cedc 100644 --- a/invokeai/backend/util/test_utils.py +++ b/invokeai/backend/util/test_utils.py @@ -7,7 +7,7 @@ import torch from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.backend.install.model_install_backend import ModelInstall -from invokeai.backend.model_management.model_manager import ModelInfo +from invokeai.backend.model_management.model_manager import LoadedModelInfo from invokeai.backend.model_management.models.base import BaseModelType, ModelNotFoundException, ModelType, SubModelType @@ -34,8 +34,8 @@ def install_and_load_model( base_model: BaseModelType, model_type: ModelType, submodel_type: Optional[SubModelType] = None, -) -> ModelInfo: - """Install a model if it is not already installed, then get the ModelInfo for that model. +) -> LoadedModelInfo: + """Install a model if it is not already installed, then get the LoadedModelInfo for that model. This is intended as a utility function for tests. @@ -49,9 +49,9 @@ def install_and_load_model( submodel_type (Optional[SubModelType]): The submodel type, forwarded to ModelManager.get_model(...). Returns: - ModelInfo + LoadedModelInfo """ - # If the requested model is already installed, return its ModelInfo. + # If the requested model is already installed, return its LoadedModelInfo. with contextlib.suppress(ModelNotFoundException): return model_installer.mgr.get_model(model_name, base_model, model_type, submodel_type) diff --git a/invokeai/invocation_api/__init__.py b/invokeai/invocation_api/__init__.py index e80bc26a00..2d3ceca11e 100644 --- a/invokeai/invocation_api/__init__.py +++ b/invokeai/invocation_api/__init__.py @@ -52,7 +52,7 @@ from invokeai.app.services.config.config_default import InvokeAIAppConfig from invokeai.app.services.image_records.image_records_common import ImageCategory from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutID -from invokeai.backend.model_management.model_manager import ModelInfo +from invokeai.backend.model_management.model_manager import LoadedModelInfo from invokeai.backend.model_management.models.base import BaseModelType, ModelType, SubModelType from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ( @@ -121,7 +121,7 @@ __all__ = [ # invokeai.app.services.config.config_default "InvokeAIAppConfig", # invokeai.backend.model_management.model_manager - "ModelInfo", + "LoadedModelInfo", # invokeai.backend.model_management.models.base "BaseModelType", "ModelType",