WIP - simplify ModelLoadRegistry

docs/contributing/MODEL_MANAGER.md

@@ -73,15 +73,6 @@ model's lifetime it may be transformed in various ways, such as
 changing its precision or converting it from a .safetensors to a
 diffusers model.
 
-`ModelType`, `ModelFormat` and `BaseModelType` are string enums that
-are defined in `invokeai.backend.model_manager.config`. They are also
-imported by, and can be reexported from,
-`invokeai.app.services.model_manager.model_records`:
-
-```
-from invokeai.app.services.model_records import ModelType, ModelFormat, BaseModelType
-```
-
 The `path` field can be absolute or relative. If relative, it is taken
 to be relative to the `models_dir` setting in the user's
 `invokeai.yaml` file.

invokeai/app/invocations/constants.py

@@ -1,6 +1,5 @@
 from typing import Literal
 
-from invokeai.backend.stable_diffusion.schedulers import SCHEDULER_MAP
 from invokeai.backend.util.devices import TorchDevice
 
 LATENT_SCALE_FACTOR = 8
@@ -11,9 +10,6 @@ factor is hard-coded to a literal '8' rather than using this constant.
 The ratio of image:latent dimensions is LATENT_SCALE_FACTOR:1, or 8:1.
 """
 
-SCHEDULER_NAME_VALUES = Literal[tuple(SCHEDULER_MAP.keys())]
-"""A literal type representing the valid scheduler names."""
-
 IMAGE_MODES = Literal["L", "RGB", "RGBA", "CMYK", "YCbCr", "LAB", "HSV", "I", "F"]
 """A literal type for PIL image modes supported by Invoke"""
 

invokeai/app/invocations/create_gradient_mask.py

@@ -19,8 +19,8 @@ from invokeai.app.invocations.fields import (
 from invokeai.app.invocations.image_to_latents import ImageToLatentsInvocation
 from invokeai.app.invocations.model import UNetField, VAEField
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.model_manager import LoadedModel
 from invokeai.backend.model_manager.config import MainConfigBase, ModelVariantType
+from invokeai.backend.model_manager.load.load_base import LoadedModel
 from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
 
 

invokeai/app/invocations/denoise_latents.py

@@ -17,7 +17,7 @@ from torchvision.transforms.functional import resize as tv_resize
 from transformers import CLIPVisionModelWithProjection
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
-from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR, SCHEDULER_NAME_VALUES
+from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
 from invokeai.app.invocations.controlnet_image_processors import ControlField
 from invokeai.app.invocations.fields import (
     ConditioningField,
@@ -53,7 +53,7 @@ from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
     TextConditioningData,
     TextConditioningRegions,
 )
-from invokeai.backend.stable_diffusion.schedulers import SCHEDULER_MAP
+from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_MAP, SCHEDULER_NAME_VALUES
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.hotfixes import ControlNetModel
 from invokeai.backend.util.mask import to_standard_float_mask

invokeai/app/invocations/image_to_latents.py

@@ -23,7 +23,7 @@ from invokeai.app.invocations.fields import (
 from invokeai.app.invocations.model import VAEField
 from invokeai.app.invocations.primitives import LatentsOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.model_manager import LoadedModel
+from invokeai.backend.model_manager.load.load_base import LoadedModel
 from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
 from invokeai.backend.stable_diffusion.vae_tiling import patch_vae_tiling_params
 

invokeai/app/invocations/scheduler.py

@@ -1,5 +1,4 @@
 from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
-from invokeai.app.invocations.constants import SCHEDULER_NAME_VALUES
 from invokeai.app.invocations.fields import (
     FieldDescriptions,
     InputField,
@@ -7,6 +6,7 @@ from invokeai.app.invocations.fields import (
     UIType,
 )
 from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES
 
 
 @invocation_output("scheduler_output")

invokeai/app/invocations/tiled_multi_diffusion_denoise_latents.py

@@ -8,7 +8,7 @@ from diffusers.schedulers.scheduling_utils import SchedulerMixin
 from pydantic import field_validator
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation
-from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR, SCHEDULER_NAME_VALUES
+from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR
 from invokeai.app.invocations.controlnet_image_processors import ControlField
 from invokeai.app.invocations.denoise_latents import DenoiseLatentsInvocation, get_scheduler
 from invokeai.app.invocations.fields import (
@@ -29,6 +29,7 @@ from invokeai.backend.stable_diffusion.multi_diffusion_pipeline import (
     MultiDiffusionPipeline,
     MultiDiffusionRegionConditioning,
 )
+from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES
 from invokeai.backend.tiles.tiles import (
     calc_tiles_min_overlap,
 )

invokeai/app/services/model_load/model_load_base.py

@@ -6,7 +6,7 @@ from pathlib import Path
 from typing import Callable, Optional
 
 from invokeai.backend.model_manager import AnyModel, AnyModelConfig, SubModelType
-from invokeai.backend.model_manager.load import LoadedModel, LoadedModelWithoutConfig
+from invokeai.backend.model_manager.load.load_base import LoadedModel, LoadedModelWithoutConfig
 from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase
 
 

invokeai/app/services/model_load/model_load_default.py

@@ -2,7 +2,7 @@
 """Implementation of model loader service."""
 
 from pathlib import Path
-from typing import Callable, Optional, Type
+from typing import Callable, Optional
 
 from picklescan.scanner import scan_file_path
 from safetensors.torch import load_file as safetensors_load_file
@@ -11,13 +11,9 @@ from torch import load as torch_load
 from invokeai.app.services.config import InvokeAIAppConfig
 from invokeai.app.services.invoker import Invoker
 from invokeai.backend.model_manager import AnyModel, AnyModelConfig, SubModelType
-from invokeai.backend.model_manager.load import (
-    LoadedModel,
-    LoadedModelWithoutConfig,
-    ModelLoaderRegistry,
-    ModelLoaderRegistryBase,
-)
+from invokeai.backend.model_manager.load.load_base import LoadedModel, LoadedModelWithoutConfig
 from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
 from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.logging import InvokeAILogger
@@ -32,7 +28,7 @@ class ModelLoadService(ModelLoadServiceBase):
         self,
         app_config: InvokeAIAppConfig,
         ram_cache: ModelCacheBase[AnyModel],
-        registry: Optional[Type[ModelLoaderRegistryBase]] = ModelLoaderRegistry,
+        registry: ModelLoaderRegistry,
     ):
         """Initialize the model load service."""
         logger = InvokeAILogger.get_logger(self.__class__.__name__)

invokeai/app/services/model_manager/__init__.py

@@ -1,17 +0,0 @@
-"""Initialization file for model manager service."""
-
-from invokeai.backend.model_manager import AnyModel, AnyModelConfig, BaseModelType, ModelType, SubModelType
-from invokeai.backend.model_manager.load import LoadedModel
-
-from .model_manager_default import ModelManagerService, ModelManagerServiceBase
-
-__all__ = [
-    "ModelManagerServiceBase",
-    "ModelManagerService",
-    "AnyModel",
-    "AnyModelConfig",
-    "BaseModelType",
-    "ModelType",
-    "SubModelType",
-    "LoadedModel",
-]

invokeai/app/services/model_manager/model_manager_default.py

@@ -7,7 +7,8 @@ import torch
 from typing_extensions import Self
 
 from invokeai.app.services.invoker import Invoker
-from invokeai.backend.model_manager.load import ModelCache, ModelLoaderRegistry
+from invokeai.backend.model_manager.load.model_cache.model_cache_default import ModelCache
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.logging import InvokeAILogger
 

invokeai/backend/ip_adapter/ip_adapter.py

@@ -11,6 +11,7 @@ from PIL import Image
 from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
 
 from invokeai.backend.ip_adapter.ip_attention_weights import IPAttentionWeights
+from invokeai.backend.model_manager.load.model_size_utils import calc_module_size
 
 from ..raw_model import RawModel
 from .resampler import Resampler
@@ -137,10 +138,7 @@ class IPAdapter(RawModel):
         self.attn_weights.to(device=self.device, dtype=self.dtype, non_blocking=non_blocking)
 
     def calc_size(self):
-        # workaround for circular import
-        from invokeai.backend.model_manager.load.model_util import calc_model_size_by_data
-
-        return calc_model_size_by_data(self._image_proj_model) + calc_model_size_by_data(self.attn_weights)
+        return calc_module_size(self._image_proj_model) + calc_module_size(self.attn_weights)
 
     def _init_image_proj_model(
         self, state_dict: dict[str, torch.Tensor]

invokeai/backend/model_manager/__init__.py

@@ -13,7 +13,6 @@ from .config import (
     SchedulerPredictionType,
     SubModelType,
 )
-from .load import LoadedModel
 from .probe import ModelProbe
 from .search import ModelSearch
 
@@ -23,7 +22,6 @@ __all__ = [
     "BaseModelType",
     "ModelRepoVariant",
     "InvalidModelConfigException",
-    "LoadedModel",
     "ModelConfigFactory",
     "ModelFormat",
     "ModelProbe",

invokeai/backend/model_manager/config.py

@@ -30,9 +30,9 @@ from diffusers.models.modeling_utils import ModelMixin
 from pydantic import BaseModel, ConfigDict, Discriminator, Field, Tag, TypeAdapter
 from typing_extensions import Annotated, Any, Dict
 
-from invokeai.app.invocations.constants import SCHEDULER_NAME_VALUES
 from invokeai.app.util.misc import uuid_string
 from invokeai.backend.model_hash.hash_validator import validate_hash
+from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES
 
 from ..raw_model import RawModel
 

invokeai/backend/model_manager/load/__init__.py

@@ -1,27 +1 @@
-# Copyright (c) 2024 Lincoln D. Stein and the InvokeAI Development Team
-"""
-Init file for the model loader.
-"""
 
-from importlib import import_module
-from pathlib import Path
-
-from .load_base import LoadedModel, LoadedModelWithoutConfig, ModelLoaderBase
-from .load_default import ModelLoader
-from .model_cache.model_cache_default import ModelCache
-from .model_loader_registry import ModelLoaderRegistry, ModelLoaderRegistryBase
-
-# This registers the subclasses that implement loaders of specific model types
-loaders = [x.stem for x in Path(Path(__file__).parent, "model_loaders").glob("*.py") if x.stem != "__init__"]
-for module in loaders:
-    import_module(f"{__package__}.model_loaders.{module}")
-
-__all__ = [
-    "LoadedModel",
-    "LoadedModelWithoutConfig",
-    "ModelCache",
-    "ModelLoaderBase",
-    "ModelLoader",
-    "ModelLoaderRegistryBase",
-    "ModelLoaderRegistry",
-]

invokeai/backend/model_manager/load/global_model_loader_registry.py

@@ -0,0 +1,8 @@
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+
+
+def _build_model_loader_registry():
+    return ModelLoaderRegistry()
+
+
+MODEL_LOADER_REGISTRY = _build_model_loader_registry()

invokeai/backend/model_manager/load/load_default.py

@@ -15,7 +15,7 @@ from invokeai.backend.model_manager import (
 from invokeai.backend.model_manager.config import DiffusersConfigBase
 from invokeai.backend.model_manager.load.load_base import LoadedModel, ModelLoaderBase
 from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase, ModelLockerBase
-from invokeai.backend.model_manager.load.model_util import calc_model_size_by_fs
+from invokeai.backend.model_manager.load.model_size_utils import calc_model_size_by_fs
 from invokeai.backend.model_manager.load.optimizations import skip_torch_weight_init
 from invokeai.backend.util.devices import TorchDevice
 

invokeai/backend/model_manager/load/model_loader_registry.py

@@ -1,48 +1,34 @@
 # 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.
+from typing import Optional, Tuple, Type
 
-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 ..config import (
-    AnyModelConfig,
-    BaseModelType,
-    ModelConfigBase,
-    ModelFormat,
-    ModelType,
-    SubModelType,
-)
-from . import ModelLoaderBase
+from invokeai.backend.model_manager.config import BaseModelType, ModelConfigBase, ModelFormat, ModelType
+from invokeai.backend.model_manager.load.load_base import AnyModelConfig, ModelLoaderBase, SubModelType
 
 
-class ModelLoaderRegistryBase(ABC):
-    """This class allows model loaders to register their type, base and format."""
+class ModelLoaderRegistry:
+    """A registry that tracks which model loader class to use for a given model type/format/base combination."""
+
+    def __init__(self):
+        self._registry: dict[str, Type[ModelLoaderBase]] = {}
 
-    @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."""
+        self,
+        loader_class: Type[ModelLoaderBase],
+        type: ModelType,
+        format: ModelFormat,
+        base: BaseModelType = BaseModelType.Any,
+    ):
+        """Register a model loader class."""
+        key = self._to_registry_key(base, type, format)
+        if key in self._registry:
+            raise RuntimeError(
+                f"{loader_class.__name__} is trying to register as a loader for {base}/{type}/{format}, but this type "
+                f"of model has already been registered by {self._registry[key].__name__}"
+            )
+        self._registry[key] = loader_class
 
-    @classmethod
-    @abstractmethod
     def get_implementation(
-        cls, config: AnyModelConfig, submodel_type: Optional[SubModelType]
+        self, config: AnyModelConfig, submodel_type: Optional[SubModelType]
     ) -> Tuple[Type[ModelLoaderBase], ModelConfigBase, Optional[SubModelType]]:
         """
         Get subclass of ModelLoaderBase registered to handle base and type.
@@ -56,46 +42,13 @@ class ModelLoaderRegistryBase(ABC):
         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)
+        key1 = self._to_registry_key(config.base, config.type, config.format)  # for a specific base type
+        key2 = self._to_registry_key(BaseModelType.Any, config.type, config.format)  # with wildcard Any
+        implementation = self._registry.get(key1, None) or self._registry.get(key2, None)
         if not implementation:
             raise NotImplementedError(
-                f"No subclass of LoadedModel is registered for base={config.base}, type={config.type}, format={config.format}"
+                f"No subclass of ModelLoaderBase is registered for base={config.base}, type={config.type}, "
+                f"format={config.format}"
             )
         return implementation, config, submodel_type
 

invokeai/backend/model_manager/load/model_loaders/generic_diffusers.py

@@ -18,8 +18,8 @@ from invokeai.backend.model_manager import (
     SubModelType,
 )
 from invokeai.backend.model_manager.config import DiffusersConfigBase
-
-from .. import ModelLoader, ModelLoaderRegistry
+from invokeai.backend.model_manager.load.load_default import ModelLoader
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
 
 
 @ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.CLIPVision, format=ModelFormat.Diffusers)

invokeai/backend/model_manager/load/model_loaders/ip_adapter.py

@@ -8,7 +8,8 @@ import torch
 
 from invokeai.backend.ip_adapter.ip_adapter import build_ip_adapter
 from invokeai.backend.model_manager import AnyModel, AnyModelConfig, BaseModelType, ModelFormat, ModelType, SubModelType
-from invokeai.backend.model_manager.load import ModelLoader, ModelLoaderRegistry
+from invokeai.backend.model_manager.load.load_default import ModelLoader
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
 from invokeai.backend.raw_model import RawModel
 
 

invokeai/backend/model_manager/load/model_size_utils.py

@@ -0,0 +1,79 @@
+import json
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+
+def calc_module_size(model: torch.nn.Module) -> int:
+    """Estimate the size of a torch.nn.Module in bytes."""
+    mem_params = sum([param.nelement() * param.element_size() for param in model.parameters()])
+    mem_bufs = sum([buf.nelement() * buf.element_size() for buf in model.buffers()])
+    mem: int = mem_params + mem_bufs  # in bytes
+    return mem
+
+
+def calc_model_size_by_fs(model_path: Path, subfolder: Optional[str] = None, variant: Optional[str] = None) -> int:
+    """Estimate the size of a model on disk in bytes."""
+    if model_path.is_file():
+        return model_path.stat().st_size
+
+    if subfolder is not None:
+        model_path = model_path / subfolder
+
+    # this can happen when, for example, the safety checker is not downloaded.
+    if not model_path.exists():
+        return 0
+
+    all_files = [f for f in model_path.iterdir() if (model_path / f).is_file()]
+
+    fp16_files = {f for f in all_files if ".fp16." in f.name or ".fp16-" in f.name}
+    bit8_files = {f for f in all_files if ".8bit." in f.name or ".8bit-" in f.name}
+    other_files = set(all_files) - fp16_files - bit8_files
+
+    if not variant:  # ModelRepoVariant.DEFAULT evaluates to empty string for compatability with HF
+        files = other_files
+    elif variant == "fp16":
+        files = fp16_files
+    elif variant == "8bit":
+        files = bit8_files
+    else:
+        raise NotImplementedError(f"Unknown variant: {variant}")
+
+    # try read from index if exists
+    index_postfix = ".index.json"
+    if variant is not None:
+        index_postfix = f".index.{variant}.json"
+
+    for file in files:
+        if not file.name.endswith(index_postfix):
+            continue
+        try:
+            with open(model_path / file, "r") as f:
+                index_data = json.loads(f.read())
+            return int(index_data["metadata"]["total_size"])
+        except Exception:
+            pass
+
+    # calculate files size if there is no index file
+    formats = [
+        (".safetensors",),  # safetensors
+        (".bin",),  # torch
+        (".onnx", ".pb"),  # onnx
+        (".msgpack",),  # flax
+        (".ckpt",),  # tf
+        (".h5",),  # tf2
+    ]
+
+    for file_format in formats:
+        model_files = [f for f in files if f.suffix in file_format]
+        if len(model_files) == 0:
+            continue
+
+        model_size = 0
+        for model_file in model_files:
+            file_stats = (model_path / model_file).stat()
+            model_size += file_stats.st_size
+        return model_size
+
+    return 0  # scheduler/feature_extractor/tokenizer - models without loading to gpu

invokeai/backend/model_manager/load/model_util.py

@@ -1,14 +1,11 @@
 # Copyright (c) 2024 The InvokeAI Development Team
 """Various utility functions needed by the loader and caching system."""
 
-import json
-from pathlib import Path
-from typing import Optional
-
 import torch
 from diffusers import DiffusionPipeline
 
 from invokeai.backend.model_manager.config import AnyModel
+from invokeai.backend.model_manager.load.model_size_utils import calc_module_size
 from invokeai.backend.onnx.onnx_runtime import IAIOnnxRuntimeModel
 
 
@@ -17,7 +14,7 @@ def calc_model_size_by_data(model: AnyModel) -> int:
     if isinstance(model, DiffusionPipeline):
         return _calc_pipeline_by_data(model)
     elif isinstance(model, torch.nn.Module):
-        return _calc_model_by_data(model)
+        return calc_module_size(model)
     elif isinstance(model, IAIOnnxRuntimeModel):
         return _calc_onnx_model_by_data(model)
     else:
@@ -30,84 +27,11 @@ def _calc_pipeline_by_data(pipeline: DiffusionPipeline) -> int:
     for submodel_key in pipeline.components.keys():
         submodel = getattr(pipeline, submodel_key)
         if submodel is not None and isinstance(submodel, torch.nn.Module):
-            res += _calc_model_by_data(submodel)
+            res += calc_module_size(submodel)
     return res
 
 
-def _calc_model_by_data(model: torch.nn.Module) -> int:
-    mem_params = sum([param.nelement() * param.element_size() for param in model.parameters()])
-    mem_bufs = sum([buf.nelement() * buf.element_size() for buf in model.buffers()])
-    mem: int = mem_params + mem_bufs  # in bytes
-    return mem
-
-
 def _calc_onnx_model_by_data(model: IAIOnnxRuntimeModel) -> int:
     tensor_size = model.tensors.size() * 2  # The session doubles this
     mem = tensor_size  # in bytes
     return mem
-
-
-def calc_model_size_by_fs(model_path: Path, subfolder: Optional[str] = None, variant: Optional[str] = None) -> int:
-    """Estimate the size of a model on disk in bytes."""
-    if model_path.is_file():
-        return model_path.stat().st_size
-
-    if subfolder is not None:
-        model_path = model_path / subfolder
-
-    # this can happen when, for example, the safety checker is not downloaded.
-    if not model_path.exists():
-        return 0
-
-    all_files = [f for f in model_path.iterdir() if (model_path / f).is_file()]
-
-    fp16_files = {f for f in all_files if ".fp16." in f.name or ".fp16-" in f.name}
-    bit8_files = {f for f in all_files if ".8bit." in f.name or ".8bit-" in f.name}
-    other_files = set(all_files) - fp16_files - bit8_files
-
-    if not variant:  # ModelRepoVariant.DEFAULT evaluates to empty string for compatability with HF
-        files = other_files
-    elif variant == "fp16":
-        files = fp16_files
-    elif variant == "8bit":
-        files = bit8_files
-    else:
-        raise NotImplementedError(f"Unknown variant: {variant}")
-
-    # try read from index if exists
-    index_postfix = ".index.json"
-    if variant is not None:
-        index_postfix = f".index.{variant}.json"
-
-    for file in files:
-        if not file.name.endswith(index_postfix):
-            continue
-        try:
-            with open(model_path / file, "r") as f:
-                index_data = json.loads(f.read())
-            return int(index_data["metadata"]["total_size"])
-        except Exception:
-            pass
-
-    # calculate files size if there is no index file
-    formats = [
-        (".safetensors",),  # safetensors
-        (".bin",),  # torch
-        (".onnx", ".pb"),  # onnx
-        (".msgpack",),  # flax
-        (".ckpt",),  # tf
-        (".h5",),  # tf2
-    ]
-
-    for file_format in formats:
-        model_files = [f for f in files if f.suffix in file_format]
-        if len(model_files) == 0:
-            continue
-
-        model_size = 0
-        for model_file in model_files:
-            file_stats = (model_path / model_file).stat()
-            model_size += file_stats.st_size
-        return model_size
-
-    return 0  # scheduler/feature_extractor/tokenizer - models without loading to gpu

invokeai/backend/stable_diffusion/schedulers/__init__.py

@@ -1,3 +0,0 @@
-from .schedulers import SCHEDULER_MAP  # noqa: F401
-
-__all__ = ["SCHEDULER_MAP"]

invokeai/backend/stable_diffusion/schedulers/schedulers.py

@@ -1,3 +1,5 @@
+from typing import Literal
+
 from diffusers import (
     DDIMScheduler,
     DDPMScheduler,
@@ -43,3 +45,9 @@ SCHEDULER_MAP = {
     "lcm": (LCMScheduler, {}),
     "tcd": (TCDScheduler, {}),
 }
+
+
+# HACK(ryand): Passing a tuple of keys to Literal works at runtime, but not at type-check time. See the docs here for
+# more info: https://typing.readthedocs.io/en/latest/spec/literal.html#parameters-at-runtime. For now, we are ignoring
+# this error. In the future, we should fix this type handling.
+SCHEDULER_NAME_VALUES = Literal[tuple(SCHEDULER_MAP.keys())]  # type: ignore

invokeai/backend/util/test_utils.py

@@ -5,9 +5,10 @@ from typing import Optional, Union
 import pytest
 import torch
 
-from invokeai.app.services.model_manager import ModelManagerServiceBase
+from invokeai.app.services.model_manager.model_manager_base import ModelManagerServiceBase
 from invokeai.app.services.model_records import UnknownModelException
-from invokeai.backend.model_manager import BaseModelType, LoadedModel, ModelType, SubModelType
+from invokeai.backend.model_manager import BaseModelType, ModelType, SubModelType
+from invokeai.backend.model_manager.load.load_base import LoadedModel
 
 
 @pytest.fixture(scope="session")

invokeai/invocation_api/__init__.py

@@ -11,7 +11,6 @@ from invokeai.app.invocations.baseinvocation import (
     invocation,
     invocation_output,
 )
-from invokeai.app.invocations.constants import SCHEDULER_NAME_VALUES
 from invokeai.app.invocations.fields import (
     BoardField,
     ColorField,
@@ -78,6 +77,7 @@ from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
     ConditioningFieldData,
     SDXLConditioningInfo,
 )
+from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES
 from invokeai.backend.util.devices import CPU_DEVICE, CUDA_DEVICE, MPS_DEVICE, choose_precision, choose_torch_device
 from invokeai.version import __version__
 
@@ -163,7 +163,7 @@ __all__ = [
     "BaseModelType",
     "ModelType",
     "SubModelType",
-    # invokeai.app.invocations.constants
+    # invokeai.backend.stable_diffusion.schedulers.schedulers
     "SCHEDULER_NAME_VALUES",
     # invokeai.version
     "__version__",

tests/app/services/model_load/test_load_api.py

@@ -5,7 +5,7 @@ import torch
 from diffusers import AutoencoderTiny
 
 from invokeai.app.services.invocation_services import InvocationServices
-from invokeai.app.services.model_manager import ModelManagerServiceBase
+from invokeai.app.services.model_manager.model_manager_base import ModelManagerServiceBase
 from invokeai.app.services.shared.invocation_context import InvocationContext, build_invocation_context
 from invokeai.backend.model_manager.load.load_base import LoadedModelWithoutConfig
 from tests.backend.model_manager.model_manager_fixtures import *  # noqa F403

tests/backend/model_manager/model_loading/test_model_load.py

@@ -4,7 +4,7 @@ Test model loading
 
 from pathlib import Path
 
-from invokeai.app.services.model_manager import ModelManagerServiceBase
+from invokeai.app.services.model_manager.model_manager_base import ModelManagerServiceBase
 from invokeai.backend.textual_inversion import TextualInversionModelRaw
 from tests.backend.model_manager.model_manager_fixtures import *  # noqa F403
 

tests/backend/model_manager/model_manager_fixtures.py

@@ -12,7 +12,8 @@ from invokeai.app.services.config import InvokeAIAppConfig
 from invokeai.app.services.download import DownloadQueueService, DownloadQueueServiceBase
 from invokeai.app.services.model_install import ModelInstallService, ModelInstallServiceBase
 from invokeai.app.services.model_load import ModelLoadService, ModelLoadServiceBase
-from invokeai.app.services.model_manager import ModelManagerService, ModelManagerServiceBase
+from invokeai.app.services.model_manager.model_manager_base import ModelManagerServiceBase
+from invokeai.app.services.model_manager.model_manager_default import ModelManagerService
 from invokeai.app.services.model_records import ModelRecordServiceBase, ModelRecordServiceSQL
 from invokeai.backend.model_manager.config import (
     BaseModelType,
@@ -25,7 +26,7 @@ from invokeai.backend.model_manager.config import (
     ModelVariantType,
     VAEDiffusersConfig,
 )
-from invokeai.backend.model_manager.load import ModelCache
+from invokeai.backend.model_manager.load.model_cache.model_cache_default import ModelCache
 from invokeai.backend.util.logging import InvokeAILogger
 from tests.backend.model_manager.model_metadata.metadata_examples import (
     HFTestLoraMetadata,