refactor(mm): migrate loaders to config classes

Author: psychedelicious

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 Any, Callable, Optional
 
 from picklescan.scanner import scan_file_path
 from safetensors.torch import load_file as safetensors_load_file
@@ -11,36 +11,29 @@
 from invokeai.app.services.config import InvokeAIAppConfig
 from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.model_load.model_load_base import ModelLoadServiceBase
-from invokeai.backend.model_manager.config import AnyModelConfig
-from invokeai.backend.model_manager.load import (
-    LoadedModel,
-    LoadedModelWithoutConfig,
-    ModelLoaderRegistry,
-    ModelLoaderRegistryBase,
-)
-from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
-from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader
+from invokeai.backend.model_manager.config import AnyModelConfig, Diffusers_Config_Base
+from invokeai.backend.model_manager.load import LoadedModel, LoadedModelWithoutConfig
+from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache, get_model_cache_key
+from invokeai.backend.model_manager.load.model_util import calc_model_size_by_fs
 from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.logging import InvokeAILogger
 
 
 class ModelLoadService(ModelLoadServiceBase):
-    """Wrapper around ModelLoaderRegistry."""
+    """Model loading service using config-based loading."""
 
     def __init__(
         self,
         app_config: InvokeAIAppConfig,
         ram_cache: ModelCache,
-        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._registry = registry
 
     def start(self, invoker: Invoker) -> None:
         self._invoker = invoker
@@ -63,18 +56,49 @@ def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubMo
         if hasattr(self, "_invoker"):
             self._invoker.services.events.emit_model_load_started(model_config, submodel_type)
 
-        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,
-        ).load_model(model_config, submodel_type)
+        loaded_model = self._load_model_from_config(model_config, submodel_type)
 
         if hasattr(self, "_invoker"):
             self._invoker.services.events.emit_model_load_complete(model_config, submodel_type)
 
         return loaded_model
 
+    def _load_model_from_config(
+        self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None
+    ) -> LoadedModel:
+        """Load a model using the config's load_model method."""
+        model_path = Path(model_config.path)
+        stats_name = ":".join([model_config.base, model_config.type, model_config.name, (submodel_type or "")])
+
+        # Check if model is already in cache
+        try:
+            cache_record = self._ram_cache.get(key=get_model_cache_key(model_config.key, submodel_type), stats_name=stats_name)
+            return LoadedModel(config=model_config, cache_record=cache_record, cache=self._ram_cache)
+        except IndexError:
+            pass
+
+        # Make room in cache
+        variant = model_config.repo_variant if isinstance(model_config, Diffusers_Config_Base) else None
+        model_size = calc_model_size_by_fs(
+            model_path=model_path,
+            subfolder=submodel_type.value if submodel_type else None,
+            variant=variant,
+        )
+        self._ram_cache.make_room(model_size)
+
+        # Load the model using the config's load_model method
+        raw_model = model_config.load_model(submodel_type)
+
+        # Cache the loaded model
+        self._ram_cache.put(
+            get_model_cache_key(model_config.key, submodel_type),
+            model=raw_model,
+        )
+
+        # Retrieve from cache and return
+        cache_record = self._ram_cache.get(key=get_model_cache_key(model_config.key, submodel_type), stats_name=stats_name)
+        return LoadedModel(config=model_config, cache_record=cache_record, cache=self._ram_cache)
+
     def load_model_from_path(
         self, model_path: Path, loader: Optional[Callable[[Path], AnyModel]] = None
     ) -> LoadedModelWithoutConfig:
@@ -107,12 +131,31 @@ def torch_load_file(checkpoint: Path) -> AnyModel:
             return result
 
         def diffusers_load_directory(directory: Path) -> AnyModel:
-            load_class = GenericDiffusersLoader(
-                app_config=self._app_config,
-                logger=self._logger,
-                ram_cache=self._ram_cache,
-                convert_cache=self.convert_cache,
-            ).get_hf_load_class(directory)
+            from diffusers.configuration_utils import ConfigMixin
+
+            class ConfigLoader(ConfigMixin):
+                """Subclass of ConfigMixin for loading diffusers configuration files."""
+
+                @classmethod
+                def load_config(cls, *args: Any, **kwargs: Any) -> dict[str, Any]:  # type: ignore
+                    """Load a diffusers ConfigMixin configuration."""
+                    cls.config_name = kwargs.pop("config_name")
+                    return super().load_config(*args, **kwargs)  # type: ignore
+
+            config = ConfigLoader.load_config(directory, config_name="config.json")
+            if class_name := config.get("_class_name"):
+                import sys
+
+                res_type = sys.modules["diffusers"]
+                load_class = getattr(res_type, class_name)
+            elif class_name := config.get("architectures"):
+                import sys
+
+                res_type = sys.modules["transformers"]
+                load_class = getattr(res_type, class_name[0])
+            else:
+                raise Exception("Unable to determine load class from config.json")
+
             return load_class.from_pretrained(model_path, torch_dtype=TorchDevice.choose_torch_dtype())
 
         loader = loader or (

invokeai/backend/flux/flux_state_dict_utils.py

@@ -1,10 +1,7 @@
-from typing import TYPE_CHECKING
+from typing import Any
 
-if TYPE_CHECKING:
-    from invokeai.backend.model_manager.legacy_probe import CkptType
 
-
-def get_flux_in_channels_from_state_dict(state_dict: "CkptType") -> int | None:
+def get_flux_in_channels_from_state_dict(state_dict: dict[str | int, Any]) -> int | None:
     """Gets the in channels from the state dict."""
 
     # "Standard" FLUX models use "img_in.weight", but some community fine tunes use

invokeai/backend/model_manager/__init__.py

@@ -6,7 +6,6 @@
     Config_Base,
     ModelConfigFactory,
 )
-from invokeai.backend.model_manager.legacy_probe import ModelProbe
 from invokeai.backend.model_manager.load import LoadedModel
 from invokeai.backend.model_manager.search import ModelSearch
 from invokeai.backend.model_manager.taxonomy import (
@@ -28,7 +27,6 @@
     "InvalidModelConfigException",
     "LoadedModel",
     "ModelConfigFactory",
-    "ModelProbe",
     "ModelSearch",
     "Config_Base",
     "AnyModel",