added textual inversion and lora loaders

invokeai/backend/model_manager/load/load_default.py

@@ -10,11 +10,17 @@ from diffusers import ModelMixin
 from diffusers.configuration_utils import ConfigMixin
 
 from invokeai.app.services.config import InvokeAIAppConfig
-from invokeai.backend.model_manager import AnyModel, AnyModelConfig, InvalidModelConfigException, ModelRepoVariant, SubModelType
+from invokeai.backend.model_manager import (
+    AnyModel,
+    AnyModelConfig,
+    InvalidModelConfigException,
+    ModelRepoVariant,
+    SubModelType,
+)
 from invokeai.backend.model_manager.load.convert_cache import ModelConvertCacheBase
 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, calc_model_size_by_data
+from invokeai.backend.model_manager.load.model_util import calc_model_size_by_data, calc_model_size_by_fs
 from invokeai.backend.model_manager.load.optimizations import skip_torch_weight_init
 from invokeai.backend.util.devices import choose_torch_device, torch_dtype
 
@@ -160,4 +166,3 @@ class ModelLoader(ModelLoaderBase):
         submodel_type: Optional[SubModelType] = None,
     ) -> AnyModel:
         raise NotImplementedError
-

invokeai/backend/model_manager/load/memory_snapshot.py

@@ -97,4 +97,4 @@ def get_pretty_snapshot_diff(snapshot_1: Optional[MemorySnapshot], snapshot_2: O
     if snapshot_1.vram is not None and snapshot_2.vram is not None:
         msg += get_msg_line("VRAM", snapshot_1.vram, snapshot_2.vram)
 
-    return "\n"+msg if len(msg)>0 else msg
+    return "\n" + msg if len(msg) > 0 else msg

invokeai/backend/model_manager/load/model_cache/__init__.py

@@ -1,5 +1,3 @@
 """Init file for RamCache."""
 
-from .model_cache_base import ModelCacheBase
-from .model_cache_default import ModelCache
 _all__ = ["ModelCacheBase", "ModelCache"]

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

@@ -14,8 +14,10 @@ from invokeai.backend.model_manager import (
 )
 from invokeai.backend.model_manager.convert_ckpt_to_diffusers import convert_controlnet_to_diffusers
 from invokeai.backend.model_manager.load.load_base import AnyModelLoader
+
 from .generic_diffusers import GenericDiffusersLoader
 
+
 @AnyModelLoader.register(base=BaseModelType.Any, type=ModelType.ControlNet, format=ModelFormat.Diffusers)
 @AnyModelLoader.register(base=BaseModelType.Any, type=ModelType.ControlNet, format=ModelFormat.Checkpoint)
 class ControlnetLoader(GenericDiffusersLoader):
@@ -37,7 +39,7 @@ class ControlnetLoader(GenericDiffusersLoader):
         if config.base not in {BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2}:
             raise Exception(f"Vae conversion not supported for model type: {config.base}")
         else:
-            assert hasattr(config, 'config')
+            assert hasattr(config, "config")
             config_file = config.config
 
         if weights_path.suffix == ".safetensors":

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

@@ -15,6 +15,7 @@ from invokeai.backend.model_manager import (
 from invokeai.backend.model_manager.load.load_base import AnyModelLoader
 from invokeai.backend.model_manager.load.load_default import ModelLoader
 
+
 @AnyModelLoader.register(base=BaseModelType.Any, type=ModelType.CLIPVision, format=ModelFormat.Diffusers)
 @AnyModelLoader.register(base=BaseModelType.Any, type=ModelType.T2IAdapter, format=ModelFormat.Diffusers)
 class GenericDiffusersLoader(ModelLoader):

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

@@ -1,11 +1,11 @@
 # Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team
 """Class for IP Adapter model loading in InvokeAI."""
 
-import torch
-
 from pathlib import Path
 from typing import Optional
 
+import torch
+
 from invokeai.backend.ip_adapter.ip_adapter import build_ip_adapter
 from invokeai.backend.model_manager import (
     AnyModel,
@@ -18,6 +18,7 @@ from invokeai.backend.model_manager import (
 from invokeai.backend.model_manager.load.load_base import AnyModelLoader
 from invokeai.backend.model_manager.load.load_default import ModelLoader
 
+
 @AnyModelLoader.register(base=BaseModelType.Any, type=ModelType.IPAdapter, format=ModelFormat.InvokeAI)
 class IPAdapterInvokeAILoader(ModelLoader):
     """Class to load IP Adapter diffusers models."""
@@ -36,4 +37,3 @@ class IPAdapterInvokeAILoader(ModelLoader):
             dtype=self._torch_dtype,
         )
         return model
-

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

@@ -2,13 +2,12 @@
 """Class for LoRA model loading in InvokeAI."""
 
 
+from logging import Logger
 from pathlib import Path
 from typing import Optional, Tuple
-from logging import Logger
 
-from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase
-from invokeai.backend.model_manager.load.convert_cache import ModelConvertCacheBase
 from invokeai.app.services.config import InvokeAIAppConfig
+from invokeai.backend.embeddings.lora import LoRAModelRaw
 from invokeai.backend.model_manager import (
     AnyModel,
     AnyModelConfig,
@@ -18,9 +17,11 @@ from invokeai.backend.model_manager import (
     ModelType,
     SubModelType,
 )
-from invokeai.backend.model_manager.lora import LoRAModelRaw
+from invokeai.backend.model_manager.load.convert_cache import ModelConvertCacheBase
 from invokeai.backend.model_manager.load.load_base import AnyModelLoader
 from invokeai.backend.model_manager.load.load_default import ModelLoader
+from invokeai.backend.model_manager.load.model_cache.model_cache_base import ModelCacheBase
+
 
 @AnyModelLoader.register(base=BaseModelType.Any, type=ModelType.Lora, format=ModelFormat.Diffusers)
 @AnyModelLoader.register(base=BaseModelType.Any, type=ModelType.Lora, format=ModelFormat.Lycoris)
@@ -47,6 +48,7 @@ class LoraLoader(ModelLoader):
     ) -> AnyModel:
         if submodel_type is not None:
             raise ValueError("There are no submodels in a LoRA model.")
+        assert self._model_base is not None
         model = LoRAModelRaw.from_checkpoint(
             file_path=model_path,
             dtype=self._torch_dtype,
@@ -56,9 +58,11 @@ class LoraLoader(ModelLoader):
 
     # override
     def _get_model_path(
-            self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None
+        self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None
     ) -> Tuple[Path, AnyModelConfig, Optional[SubModelType]]:
-        self._model_base = config.base  # cheating a little - setting this variable for later call to _load_model()
+        self._model_base = (
+            config.base
+        )  # cheating a little - we remember this variable for using in the subsequent call to _load_model()
 
         model_base_path = self._app_config.models_path
         model_path = model_base_path / config.path
@@ -72,5 +76,3 @@ class LoraLoader(ModelLoader):
 
         result = model_path.resolve(), config, submodel_type
         return result
-
-

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

@@ -0,0 +1,55 @@
+# Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team
+"""Class for TI model loading in InvokeAI."""
+
+
+from pathlib import Path
+from typing import Optional, Tuple
+
+from invokeai.backend.embeddings.model_patcher import TextualInversionModel as TextualInversionModelRaw
+from invokeai.backend.model_manager import (
+    AnyModel,
+    AnyModelConfig,
+    BaseModelType,
+    ModelFormat,
+    ModelRepoVariant,
+    ModelType,
+    SubModelType,
+)
+from invokeai.backend.model_manager.load.load_base import AnyModelLoader
+from invokeai.backend.model_manager.load.load_default import ModelLoader
+
+
+@AnyModelLoader.register(base=BaseModelType.Any, type=ModelType.TextualInversion, format=ModelFormat.EmbeddingFile)
+@AnyModelLoader.register(base=BaseModelType.Any, type=ModelType.TextualInversion, format=ModelFormat.EmbeddingFolder)
+class TextualInversionLoader(ModelLoader):
+    """Class to load TI models."""
+
+    def _load_model(
+        self,
+        model_path: Path,
+        model_variant: Optional[ModelRepoVariant] = None,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if submodel_type is not None:
+            raise ValueError("There are no submodels in a TI model.")
+        model = TextualInversionModelRaw.from_checkpoint(
+            file_path=model_path,
+            dtype=self._torch_dtype,
+        )
+        return model
+
+    # override
+    def _get_model_path(
+        self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None
+    ) -> Tuple[Path, AnyModelConfig, Optional[SubModelType]]:
+        model_path = self._app_config.models_path / config.path
+
+        if config.format == ModelFormat.EmbeddingFolder:
+            path = model_path / "learned_embeds.bin"
+        else:
+            path = model_path
+
+        if not path.exists():
+            raise OSError(f"The embedding file at {path} was not found")
+
+        return path, config, submodel_type

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

@@ -15,6 +15,7 @@ from invokeai.backend.model_manager import (
 )
 from invokeai.backend.model_manager.convert_ckpt_to_diffusers import convert_ldm_vae_to_diffusers
 from invokeai.backend.model_manager.load.load_base import AnyModelLoader
+
 from .generic_diffusers import GenericDiffusersLoader
 
 

invokeai/backend/model_manager/load/model_util.py

@@ -3,13 +3,13 @@
 
 import json
 from pathlib import Path
-from typing import Optional, Union
+from typing import Optional
 
 import torch
 from diffusers import DiffusionPipeline
 
 from invokeai.backend.model_manager.config import AnyModel
-from invokeai.backend.model_manager.onnx_runtime import IAIOnnxRuntimeModel
+from invokeai.backend.onnx.onnx_runtime import IAIOnnxRuntimeModel
 
 
 def calc_model_size_by_data(model: AnyModel) -> int: