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.
This commit is contained in:
psychedelicious
2024-02-11 09:27:57 +11:00
parent e0694a2856
commit 0f8af643d1
9 changed files with 38 additions and 30 deletions

View File

@ -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

View File

@ -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

View File

@ -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 = "<NO_HASH>" # TODO:
return ModelInfo(
return LoadedModelInfo(
context=model_context,
name=model_name,
base_model=base_model,

View File

@ -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)