mirror of
https://github.com/invoke-ai/InvokeAI.git
synced 2025-04-04 22:43:40 +08:00
105 lines
3.9 KiB
Python
105 lines
3.9 KiB
Python
# Copyright (c) 2024 Lincoln D. Stein and the InvokeAI Development team
|
|
"""
|
|
This module implements a system in which model loaders register the
|
|
type, base and format of models that they know how to load.
|
|
|
|
Use like this:
|
|
|
|
cls, model_config, submodel_type = ModelLoaderRegistry.get_implementation(model_config, submodel_type) # type: ignore
|
|
loaded_model = cls(
|
|
app_config=app_config,
|
|
logger=logger,
|
|
ram_cache=ram_cache,
|
|
convert_cache=convert_cache
|
|
).load_model(model_config, submodel_type)
|
|
|
|
"""
|
|
|
|
from abc import ABC, abstractmethod
|
|
from typing import Callable, Dict, Optional, Tuple, Type, TypeVar
|
|
|
|
from invokeai.backend.model_manager.config import (
|
|
AnyModelConfig,
|
|
BaseModelType,
|
|
ModelConfigBase,
|
|
ModelFormat,
|
|
ModelType,
|
|
SubModelType,
|
|
)
|
|
from invokeai.backend.model_manager.load import ModelLoaderBase
|
|
|
|
|
|
class ModelLoaderRegistryBase(ABC):
|
|
"""This class allows model loaders to register their type, base and format."""
|
|
|
|
@classmethod
|
|
@abstractmethod
|
|
def register(
|
|
cls, type: ModelType, format: ModelFormat, base: BaseModelType = BaseModelType.Any
|
|
) -> Callable[[Type[ModelLoaderBase]], Type[ModelLoaderBase]]:
|
|
"""Define a decorator which registers the subclass of loader."""
|
|
|
|
@classmethod
|
|
@abstractmethod
|
|
def get_implementation(
|
|
cls, config: AnyModelConfig, submodel_type: Optional[SubModelType]
|
|
) -> Tuple[Type[ModelLoaderBase], ModelConfigBase, Optional[SubModelType]]:
|
|
"""
|
|
Get subclass of ModelLoaderBase registered to handle base and type.
|
|
|
|
Parameters:
|
|
:param config: Model configuration record, as returned by ModelRecordService
|
|
:param submodel_type: Submodel to fetch (main models only)
|
|
:return: tuple(loader_class, model_config, submodel_type)
|
|
|
|
Note that the returned model config may be different from one what passed
|
|
in, in the event that a submodel type is provided.
|
|
"""
|
|
|
|
|
|
TModelLoader = TypeVar("TModelLoader", bound=ModelLoaderBase)
|
|
|
|
|
|
class ModelLoaderRegistry(ModelLoaderRegistryBase):
|
|
"""
|
|
This class allows model loaders to register their type, base and format.
|
|
"""
|
|
|
|
_registry: Dict[str, Type[ModelLoaderBase]] = {}
|
|
|
|
@classmethod
|
|
def register(
|
|
cls, type: ModelType, format: ModelFormat, base: BaseModelType = BaseModelType.Any
|
|
) -> Callable[[Type[TModelLoader]], Type[TModelLoader]]:
|
|
"""Define a decorator which registers the subclass of loader."""
|
|
|
|
def decorator(subclass: Type[TModelLoader]) -> Type[TModelLoader]:
|
|
key = cls._to_registry_key(base, type, format)
|
|
if key in cls._registry:
|
|
raise Exception(
|
|
f"{subclass.__name__} is trying to register as a loader for {base}/{type}/{format}, but this type of model has already been registered by {cls._registry[key].__name__}"
|
|
)
|
|
cls._registry[key] = subclass
|
|
return subclass
|
|
|
|
return decorator
|
|
|
|
@classmethod
|
|
def get_implementation(
|
|
cls, config: AnyModelConfig, submodel_type: Optional[SubModelType]
|
|
) -> Tuple[Type[ModelLoaderBase], ModelConfigBase, Optional[SubModelType]]:
|
|
"""Get subclass of ModelLoaderBase registered to handle base and type."""
|
|
|
|
key1 = cls._to_registry_key(config.base, config.type, config.format) # for a specific base type
|
|
key2 = cls._to_registry_key(BaseModelType.Any, config.type, config.format) # with wildcard Any
|
|
implementation = cls._registry.get(key1) or cls._registry.get(key2)
|
|
if not implementation:
|
|
raise NotImplementedError(
|
|
f"No subclass of LoadedModel is registered for base={config.base}, type={config.type}, format={config.format}"
|
|
)
|
|
return implementation, config, submodel_type
|
|
|
|
@staticmethod
|
|
def _to_registry_key(base: BaseModelType, type: ModelType, format: ModelFormat) -> str:
|
|
return "-".join([base.value, type.value, format.value])
|