mirror of
https://github.com/invoke-ai/InvokeAI.git
synced 2025-04-04 22:43:40 +08:00
5a3195f757
- Replace AnyModelLoader with ModelLoaderRegistry - Fix type check errors in multiple files - Remove apparently unneeded `get_model_config_enum()` method from model manager - Remove last vestiges of old model manager - Updated tests and documentation resolve conflict with seamless.py
117 lines
4.3 KiB
Python
117 lines
4.3 KiB
Python
# Copyright (c) 2024 Lincoln D. Stein and the InvokeAI Team
|
|
"""Implementation of model loader service."""
|
|
|
|
from typing import Optional, Type
|
|
|
|
from invokeai.app.services.config import InvokeAIAppConfig
|
|
from invokeai.app.services.invocation_processor.invocation_processor_common import CanceledException
|
|
from invokeai.app.services.invoker import Invoker
|
|
from invokeai.app.services.shared.invocation_context import InvocationContextData
|
|
from invokeai.backend.model_manager import AnyModel, AnyModelConfig, SubModelType
|
|
from invokeai.backend.model_manager.load import (
|
|
LoadedModel,
|
|
ModelLoaderRegistry,
|
|
ModelLoaderRegistryBase,
|
|
)
|
|
from invokeai.backend.model_manager.load.convert_cache import ModelConvertCacheBase
|
|
from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase
|
|
from invokeai.backend.util.logging import InvokeAILogger
|
|
|
|
from .model_load_base import ModelLoadServiceBase
|
|
|
|
|
|
class ModelLoadService(ModelLoadServiceBase):
|
|
"""Wrapper around ModelLoaderRegistry."""
|
|
|
|
def __init__(
|
|
self,
|
|
app_config: InvokeAIAppConfig,
|
|
ram_cache: ModelCacheBase[AnyModel],
|
|
convert_cache: ModelConvertCacheBase,
|
|
registry: Optional[Type[ModelLoaderRegistryBase]] = ModelLoaderRegistry,
|
|
):
|
|
"""Initialize the model load service."""
|
|
logger = InvokeAILogger.get_logger(self.__class__.__name__)
|
|
logger.setLevel(app_config.log_level.upper())
|
|
self._logger = logger
|
|
self._app_config = app_config
|
|
self._ram_cache = ram_cache
|
|
self._convert_cache = convert_cache
|
|
self._registry = registry
|
|
|
|
def start(self, invoker: Invoker) -> None:
|
|
self._invoker = invoker
|
|
|
|
@property
|
|
def ram_cache(self) -> ModelCacheBase[AnyModel]:
|
|
"""Return the RAM cache used by this loader."""
|
|
return self._ram_cache
|
|
|
|
@property
|
|
def convert_cache(self) -> ModelConvertCacheBase:
|
|
"""Return the checkpoint convert cache used by this loader."""
|
|
return self._convert_cache
|
|
|
|
def load_model(
|
|
self,
|
|
model_config: AnyModelConfig,
|
|
submodel_type: Optional[SubModelType] = None,
|
|
context_data: Optional[InvocationContextData] = None,
|
|
) -> LoadedModel:
|
|
"""
|
|
Given a model's configuration, load it and return the LoadedModel object.
|
|
|
|
:param model_config: Model configuration record (as returned by ModelRecordBase.get_model())
|
|
:param submodel: For main (pipeline models), the submodel to fetch.
|
|
:param context: Invocation context used for event reporting
|
|
"""
|
|
if context_data:
|
|
self._emit_load_event(
|
|
context_data=context_data,
|
|
model_config=model_config,
|
|
)
|
|
|
|
implementation, model_config, submodel_type = self._registry.get_implementation(model_config, submodel_type) # type: ignore
|
|
loaded_model: LoadedModel = implementation(
|
|
app_config=self._app_config,
|
|
logger=self._logger,
|
|
ram_cache=self._ram_cache,
|
|
convert_cache=self._convert_cache,
|
|
).load_model(model_config, submodel_type)
|
|
|
|
if context_data:
|
|
self._emit_load_event(
|
|
context_data=context_data,
|
|
model_config=model_config,
|
|
loaded=True,
|
|
)
|
|
return loaded_model
|
|
|
|
def _emit_load_event(
|
|
self,
|
|
context_data: InvocationContextData,
|
|
model_config: AnyModelConfig,
|
|
loaded: Optional[bool] = False,
|
|
) -> None:
|
|
if not self._invoker:
|
|
return
|
|
if self._invoker.services.queue.is_canceled(context_data.session_id):
|
|
raise CanceledException()
|
|
|
|
if not loaded:
|
|
self._invoker.services.events.emit_model_load_started(
|
|
queue_id=context_data.queue_id,
|
|
queue_item_id=context_data.queue_item_id,
|
|
queue_batch_id=context_data.batch_id,
|
|
graph_execution_state_id=context_data.session_id,
|
|
model_config=model_config,
|
|
)
|
|
else:
|
|
self._invoker.services.events.emit_model_load_completed(
|
|
queue_id=context_data.queue_id,
|
|
queue_item_id=context_data.queue_item_id,
|
|
queue_batch_id=context_data.batch_id,
|
|
graph_execution_state_id=context_data.session_id,
|
|
model_config=model_config,
|
|
)
|