move ModelManager initialization into its own module and restore embedding support

invokeai/app/api/dependencies.py

@@ -4,7 +4,7 @@ import os
 from argparse import Namespace
 
 from ...backend import Globals
-from ..services.generate_initializer import get_model_manager
+from ..services.model_manager_initializer import get_model_manager
 from ..services.graph import GraphExecutionState
 from ..services.image_storage import DiskImageStorage
 from ..services.invocation_queue import MemoryInvocationQueue
@@ -47,8 +47,6 @@ class ApiDependencies:
         # TODO: Use a logger
         print(f">> Internet connectivity is {Globals.internet_available}")
 
-        model_manager = get_model_manager(args, config)
-
         events = FastAPIEventService(event_handler_id)
 
         output_folder = os.path.abspath(
@@ -61,7 +59,7 @@ class ApiDependencies:
         db_location = os.path.join(output_folder, "invokeai.db")
 
         services = InvocationServices(
-            generator_factory=generator_factory,
+            model_manager=get_model_manager(args, config),
             events=events,
             images=images,
             queue=MemoryInvocationQueue(),

invokeai/app/services/generate_initializer.py

@@ -6,97 +6,8 @@ from argparse import Namespace
 from omegaconf import OmegaConf
 
 import invokeai.version
-from ...backend import ModelManager
-from ...backend.util import choose_precision, choose_torch_device
 from ...backend import Globals
 
-# TODO: most of this code should be split into individual services as the Generate.py code is deprecated
-def get_model_manager(args, config) -> ModelManager:
-    if not args.conf:
-        config_file = os.path.join(Globals.root, "configs", "models.yaml")
-        if not os.path.exists(config_file):
-            report_model_error(
-                args, FileNotFoundError(f"The file {config_file} could not be found.")
-            )
-
-    print(f">> {invokeai.version.__app_name__}, version {invokeai.version.__version__}")
-    print(f'>> InvokeAI runtime directory is "{Globals.root}"')
-
-    # these two lines prevent a horrible warning message from appearing
-    # when the frozen CLIP tokenizer is imported
-    import transformers  # type: ignore
-
-    transformers.logging.set_verbosity_error()
-    import diffusers
-
-    diffusers.logging.set_verbosity_error()
-
-    # Loading Face Restoration and ESRGAN Modules
-    gfpgan, codeformer, esrgan = load_face_restoration(args)
-
-    # normalize the config directory relative to root
-    if not os.path.isabs(args.conf):
-        args.conf = os.path.normpath(os.path.join(Globals.root, args.conf))
-
-    if args.embeddings:
-        if not os.path.isabs(args.embedding_path):
-            embedding_path = os.path.normpath(
-                os.path.join(Globals.root, args.embedding_path)
-            )
-        else:
-            embedding_path = args.embedding_path
-    else:
-        embedding_path = None
-
-    # migrate legacy models
-    ModelManager.migrate_models()
-
-    # load the infile as a list of lines
-    if args.infile:
-        try:
-            if os.path.isfile(args.infile):
-                infile = open(args.infile, "r", encoding="utf-8")
-            elif args.infile == "-":  # stdin
-                infile = sys.stdin
-            else:
-                raise FileNotFoundError(f"{args.infile} not found.")
-        except (FileNotFoundError, IOError) as e:
-            print(f"{e}. Aborting.")
-            sys.exit(-1)
-
-    # creating the model manager
-    try:
-        device = torch.device(choose_torch_device())
-        precision = 'float16' if args.precision=='float16' \
-        else 'float32' if args.precision=='float32' \
-        else choose_precision(device)
-        
-        model_manager = ModelManager(
-            OmegaConf.load(args.conf),
-            precision=precision,
-            device_type=device,
-            max_loaded_models=args.max_loaded_models,
-        )
-    except (FileNotFoundError, TypeError, AssertionError) as e:
-        report_model_error(args, e)
-    except (IOError, KeyError) as e:
-        print(f"{e}. Aborting.")
-        sys.exit(-1)
-
-    if args.seamless:
-        #TODO: do something here ?
-        print(">> changed to seamless tiling mode")
-
-    # try to autoconvert new models
-    # autoimport new .ckpt files
-    if path := args.autoconvert:
-        model_manager.autoconvert_weights(
-            conf_path=args.conf,
-            weights_directory=path,
-        )
-
-    return model_manager
-
 def load_face_restoration(opt):
     try:
         gfpgan, codeformer, esrgan = None, None, None
@@ -122,42 +33,3 @@ def load_face_restoration(opt):
     return gfpgan, codeformer, esrgan
 
 
-def report_model_error(opt: Namespace, e: Exception):
-    print(f'** An error occurred while attempting to initialize the model: "{str(e)}"')
-    print(
-        "** This can be caused by a missing or corrupted models file, and can sometimes be fixed by (re)installing the models."
-    )
-    yes_to_all = os.environ.get("INVOKE_MODEL_RECONFIGURE")
-    if yes_to_all:
-        print(
-            "** Reconfiguration is being forced by environment variable INVOKE_MODEL_RECONFIGURE"
-        )
-    else:
-        response = input(
-            "Do you want to run invokeai-configure script to select and/or reinstall models? [y] "
-        )
-        if response.startswith(("n", "N")):
-            return
-
-    print("invokeai-configure is launching....\n")
-
-    # Match arguments that were set on the CLI
-    # only the arguments accepted by the configuration script are parsed
-    root_dir = ["--root", opt.root_dir] if opt.root_dir is not None else []
-    config = ["--config", opt.conf] if opt.conf is not None else []
-    previous_args = sys.argv
-    sys.argv = ["invokeai-configure"]
-    sys.argv.extend(root_dir)
-    sys.argv.extend(config)
-    if yes_to_all is not None:
-        for arg in yes_to_all.split():
-            sys.argv.append(arg)
-
-    from invokeai.frontend.install import invokeai_configure
-
-    invokeai_configure()
-    # TODO: Figure out how to restart
-    # print('** InvokeAI will now restart')
-    # sys.argv = previous_args
-    # main() # would rather do a os.exec(), but doesn't exist?
-    # sys.exit(0)

invokeai/app/services/model_manager_initializer.py

@@ -0,0 +1,136 @@
+import os
+import sys
+import torch
+from argparse import Namespace
+from omegaconf import OmegaConf
+from pathlib import Path
+
+import invokeai.version
+from ...backend import ModelManager
+from ...backend.util import choose_precision, choose_torch_device
+from ...backend import Globals
+
+# TODO: most of this code should be split into individual services as the Generate.py code is deprecated
+def get_model_manager(args, config) -> ModelManager:
+    if not args.conf:
+        config_file = os.path.join(Globals.root, "configs", "models.yaml")
+        if not os.path.exists(config_file):
+            report_model_error(
+                args, FileNotFoundError(f"The file {config_file} could not be found.")
+            )
+
+    print(f">> {invokeai.version.__app_name__}, version {invokeai.version.__version__}")
+    print(f'>> InvokeAI runtime directory is "{Globals.root}"')
+
+    # these two lines prevent a horrible warning message from appearing
+    # when the frozen CLIP tokenizer is imported
+    import transformers  # type: ignore
+
+    transformers.logging.set_verbosity_error()
+    import diffusers
+
+    diffusers.logging.set_verbosity_error()
+
+    # normalize the config directory relative to root
+    if not os.path.isabs(args.conf):
+        args.conf = os.path.normpath(os.path.join(Globals.root, args.conf))
+
+    if args.embeddings:
+        if not os.path.isabs(args.embedding_path):
+            embedding_path = os.path.normpath(
+                os.path.join(Globals.root, args.embedding_path)
+            )
+        else:
+            embedding_path = args.embedding_path
+    else:
+        embedding_path = None
+
+    # migrate legacy models
+    ModelManager.migrate_models()
+
+    # load the infile as a list of lines
+    if args.infile:
+        try:
+            if os.path.isfile(args.infile):
+                infile = open(args.infile, "r", encoding="utf-8")
+            elif args.infile == "-":  # stdin
+                infile = sys.stdin
+            else:
+                raise FileNotFoundError(f"{args.infile} not found.")
+        except (FileNotFoundError, IOError) as e:
+            print(f"{e}. Aborting.")
+            sys.exit(-1)
+
+    # creating the model manager
+    try:
+        device = torch.device(choose_torch_device())
+        precision = 'float16' if args.precision=='float16' \
+        else 'float32' if args.precision=='float32' \
+        else choose_precision(device)
+        
+        model_manager = ModelManager(
+            OmegaConf.load(args.conf),
+            precision=precision,
+            device_type=device,
+            max_loaded_models=args.max_loaded_models,
+            embedding_path = Path(embedding_path),
+        )
+    except (FileNotFoundError, TypeError, AssertionError) as e:
+        report_model_error(args, e)
+    except (IOError, KeyError) as e:
+        print(f"{e}. Aborting.")
+        sys.exit(-1)
+
+    if args.seamless:
+        #TODO: do something here ?
+        print(">> changed to seamless tiling mode")
+
+    # try to autoconvert new models
+    # autoimport new .ckpt files
+    if path := args.autoconvert:
+        model_manager.autoconvert_weights(
+            conf_path=args.conf,
+            weights_directory=path,
+        )
+
+    return model_manager
+
+def report_model_error(opt: Namespace, e: Exception):
+    print(f'** An error occurred while attempting to initialize the model: "{str(e)}"')
+    print(
+        "** This can be caused by a missing or corrupted models file, and can sometimes be fixed by (re)installing the models."
+    )
+    yes_to_all = os.environ.get("INVOKE_MODEL_RECONFIGURE")
+    if yes_to_all:
+        print(
+            "** Reconfiguration is being forced by environment variable INVOKE_MODEL_RECONFIGURE"
+        )
+    else:
+        response = input(
+            "Do you want to run invokeai-configure script to select and/or reinstall models? [y] "
+        )
+        if response.startswith(("n", "N")):
+            return
+
+    print("invokeai-configure is launching....\n")
+
+    # Match arguments that were set on the CLI
+    # only the arguments accepted by the configuration script are parsed
+    root_dir = ["--root", opt.root_dir] if opt.root_dir is not None else []
+    config = ["--config", opt.conf] if opt.conf is not None else []
+    previous_args = sys.argv
+    sys.argv = ["invokeai-configure"]
+    sys.argv.extend(root_dir)
+    sys.argv.extend(config)
+    if yes_to_all is not None:
+        for arg in yes_to_all.split():
+            sys.argv.append(arg)
+
+    from invokeai.frontend.install import invokeai_configure
+
+    invokeai_configure()
+    # TODO: Figure out how to restart
+    # print('** InvokeAI will now restart')
+    # sys.argv = previous_args
+    # main() # would rather do a os.exec(), but doesn't exist?
+    # sys.exit(0)

invokeai/backend/generate.py

@@ -222,6 +222,7 @@ class Generate:
             self.precision,
             max_loaded_models=max_loaded_models,
             sequential_offload=self.free_gpu_mem,
+            embedding_path=Path(self.embedding_path),
         )
         # don't accept invalid models
         fallback = self.model_manager.default_model() or FALLBACK_MODEL_NAME
@@ -940,18 +941,6 @@ class Generate:
         self.generators = {}
 
         set_seed(random.randrange(0, np.iinfo(np.uint32).max))
-        if self.embedding_path is not None:
-            print(f">> Loading embeddings from {self.embedding_path}")
-            for root, _, files in os.walk(self.embedding_path):
-                for name in files:
-                    ti_path = os.path.join(root, name)
-                    self.model.textual_inversion_manager.load_textual_inversion(
-                        ti_path, defer_injecting_tokens=True
-                    )
-            print(
-                f'>> Textual inversion triggers: {", ".join(sorted(self.model.textual_inversion_manager.get_all_trigger_strings()))}'
-            )
-
         self.model_name = model_name
         self._set_scheduler()  # requires self.model_name to be set first
         return self.model

invokeai/backend/model_management/model_manager.py

@@ -54,12 +54,13 @@ class ModelManager(object):
     Model manager handles loading, caching, importing, deleting, converting, and editing models.
     '''
     def __init__(
-        self,
+            self,
-        config: OmegaConf|Path,
+            config: OmegaConf|Path,
-        device_type: torch.device = CUDA_DEVICE,
+            device_type: torch.device = CUDA_DEVICE,
-        precision: str = "float16",
+            precision: str = "float16",
-        max_loaded_models=DEFAULT_MAX_MODELS,
+            max_loaded_models=DEFAULT_MAX_MODELS,
-        sequential_offload=False,
+            sequential_offload=False,
+            embedding_path: Path=None,
     ):
         """
         Initialize with the path to the models.yaml config file or
@@ -80,6 +81,7 @@ class ModelManager(object):
         self.stack = []  # this is an LRU FIFO
         self.current_model = None
         self.sequential_offload = sequential_offload
+        self.embedding_path = embedding_path
 
     def valid_model(self, model_name: str) -> bool:
         """
@@ -434,6 +436,7 @@ class ModelManager(object):
         height = width
 
         print(f"  | Default image dimensions = {width} x {height}")
+        self._add_embeddings_to_model(pipeline)
 
         return pipeline, width, height, model_hash
 
@@ -1070,6 +1073,19 @@ class ModelManager(object):
             self.stack.remove(model_name)
         self.stack.append(model_name)
 
+    def _add_embeddings_to_model(self, model: StableDiffusionGeneratorPipeline):
+        if self.embedding_path is not None:
+            print(f">> Loading embeddings from {self.embedding_path}")
+            for root, _, files in os.walk(self.embedding_path):
+                for name in files:
+                    ti_path = os.path.join(root, name)
+                    model.textual_inversion_manager.load_textual_inversion(
+                        ti_path, defer_injecting_tokens=True
+                    )
+            print(
+                f'>> Textual inversion triggers: {", ".join(sorted(model.textual_inversion_manager.get_all_trigger_strings()))}'
+            )
+            
     def _has_cuda(self) -> bool:
         return self.device.type == "cuda"