resolved conflicts; ran black and isort

ldm/invoke/__init__.py

@@ -2,3 +2,12 @@ from ._version import __version__
 
 __app_id__= 'invoke-ai/InvokeAI'
 __app_name__= 'InvokeAI'
+
+
+def _ignore_xformers_triton_message_on_windows():
+    import logging
+    logging.getLogger("xformers").addFilter(
+        lambda record: 'A matching Triton is not available' not in record.getMessage())
+
+# In order to be effective, this needs to happen before anything could possibly import xformers.
+_ignore_xformers_triton_message_on_windows()

ldm/invoke/args.py

@@ -758,6 +758,9 @@ class Args(object):
                 !fix applies upscaling/facefixing to a previously-generated image.
                 invoke> !fix 0000045.4829112.png -G1 -U4 -ft codeformer
 
+            *embeddings*
+                invoke> !triggers     -- return all trigger phrases contained in loaded embedding files
+
             *History manipulation*
             !fetch retrieves the command used to generate an earlier image. Provide
             a directory wildcard and the name of a file to write and all the commands

ldm/invoke/ckpt_to_diffuser.py

@@ -53,6 +53,7 @@ from diffusers import (
 )
 from diffusers.pipelines.latent_diffusion.pipeline_latent_diffusion import LDMBertConfig, LDMBertModel
 from diffusers.pipelines.paint_by_example import PaintByExampleImageEncoder, PaintByExamplePipeline
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
 from diffusers.utils import is_safetensors_available
 from transformers import AutoFeatureExtractor, BertTokenizerFast, CLIPTextModel, CLIPTokenizer, CLIPVisionConfig
 
@@ -984,6 +985,7 @@ def load_pipeline_from_original_stable_diffusion_ckpt(
         elif model_type in ['FrozenCLIPEmbedder','WeightedFrozenCLIPEmbedder']:
             text_model = convert_ldm_clip_checkpoint(checkpoint)
             tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14",cache_dir=cache_dir)
+            safety_checker = StableDiffusionSafetyChecker.from_pretrained('CompVis/stable-diffusion-safety-checker',cache_dir=global_cache_dir("hub"))
             feature_extractor = AutoFeatureExtractor.from_pretrained("CompVis/stable-diffusion-safety-checker",cache_dir=cache_dir)
             pipe = pipeline_class(
                 vae=vae,
@@ -991,7 +993,7 @@ def load_pipeline_from_original_stable_diffusion_ckpt(
                 tokenizer=tokenizer,
                 unet=unet,
                 scheduler=scheduler,
-                safety_checker=None,
+                safety_checker=safety_checker,
                 feature_extractor=feature_extractor,
             )
         else:

ldm/invoke/conditioning.py

@@ -93,7 +93,7 @@ def _get_conditioning_for_prompt(parsed_prompt: Union[Blend, FlattenedPrompt], p
     Process prompt structure and tokens, and return (conditioning, unconditioning, extra_conditioning_info)
     """
 
-    if log_tokens or Globals.log_tokenization:
+    if log_tokens or getattr(Globals, "log_tokenization", False):
         print(f"\n>> [TOKENLOG] Parsed Prompt: {parsed_prompt}")
         print(f"\n>> [TOKENLOG] Parsed Negative Prompt: {parsed_negative_prompt}")
 
@@ -236,7 +236,7 @@ def _get_embeddings_and_tokens_for_prompt(model, flattened_prompt: FlattenedProm
     fragments = [x.text for x in flattened_prompt.children]
     weights = [x.weight for x in flattened_prompt.children]
     embeddings, tokens = model.get_learned_conditioning([fragments], return_tokens=True, fragment_weights=[weights])
-    if log_tokens or Globals.log_tokenization:
+    if log_tokens or getattr(Globals, "log_tokenization", False):
         text = " ".join(fragments)
         log_tokenization(text, model, display_label=log_display_label)
 
@@ -296,4 +296,4 @@ def log_tokenization(text, model, display_label=None):
     
     if discarded != "":
         print(f'\n>> [TOKENLOG] Tokens Discarded ({totalTokens - usedTokens}):')
-        print(f'{discarded}\x1b[0m')
+        print(f'{discarded}\x1b[0m')

ldm/invoke/config/invokeai_configure.py

@@ -9,7 +9,6 @@
 print("Loading Python libraries...\n")
 import argparse
 import curses
-import npyscreen
 import io
 import os
 import re
@@ -17,14 +16,14 @@ import shutil
 import sys
 import traceback
 import warnings
+from argparse import Namespace
 from pathlib import Path
 from urllib import request
 
+import npyscreen
 import transformers
-from argparse import Namespace
 from huggingface_hub import HfFolder
 from huggingface_hub import login as hf_hub_login
-
 from omegaconf import OmegaConf
 from tqdm import tqdm
 from transformers import (
@@ -35,16 +34,17 @@ from transformers import (
 )
 
 import invokeai.configs as configs
-from ..args import Args, PRECISION_CHOICES
-from .model_install_backend import (
-    download_from_hf,
-    recommended_datasets,
-    default_dataset,
-)
-from .model_install import process_and_execute, addModelsForm
-from .widgets import IntTitleSlider
+
+from ..args import PRECISION_CHOICES, Args
 from ..globals import Globals, global_config_dir
 from ..readline import generic_completer
+from .model_install import addModelsForm, process_and_execute
+from .model_install_backend import (
+    default_dataset,
+    download_from_hf,
+    recommended_datasets,
+)
+from .widgets import IntTitleSlider
 
 warnings.filterwarnings("ignore")
 
@@ -500,7 +500,7 @@ class editOptsForm(npyscreen.FormMultiPage):
         options = self.marshall_arguments()
         if self.validate_field_values(options):
             self.parentApp.new_opts = options
-            if hasattr(self.parentApp,'model_select'):
+            if hasattr(self.parentApp, "model_select"):
                 self.parentApp.setNextForm("MODELS")
             else:
                 self.parentApp.setNextForm(None)
@@ -578,18 +578,21 @@ class EditOptApplication(npyscreen.NPSAppManaged):
     def new_opts(self):
         return self.options.marshall_arguments()
 
+
 def edit_opts(program_opts: Namespace, invokeai_opts: Namespace) -> argparse.Namespace:
     editApp = EditOptApplication(program_opts, invokeai_opts)
     editApp.run()
     return editApp.new_opts()
 
-def default_startup_options()->Namespace:
+
+def default_startup_options() -> Namespace:
     opts = Args().parse_args([])
     opts.outdir = str(default_output_dir())
     opts.safety_checker = True
     return opts
 
-def default_user_selections(program_opts: Namespace)->Namespace:
+
+def default_user_selections(program_opts: Namespace) -> Namespace:
     return Namespace(
         starter_models=recommended_datasets()
         if program_opts.yes_to_all
@@ -602,6 +605,8 @@ def default_user_selections(program_opts: Namespace)->Namespace:
         import_model_paths=None,
         convert_to_diffusers=None,
     )
+
+
 # -------------------------------------
 def initialize_rootdir(root: str, yes_to_all: bool = False):
     print("** INITIALIZING INVOKEAI RUNTIME DIRECTORY **")
@@ -758,7 +763,7 @@ def main():
 
     try:
         models_to_download = default_user_selections(opt)
-        
+
         # We check for to see if the runtime directory is correctly initialized.
         init_file = Path(Globals.root, Globals.initfile)
         if not init_file.exists():
@@ -771,7 +776,9 @@ def main():
             if init_options:
                 write_opts(init_options, init_file)
             else:
-                print("\n** CANCELLED AT USER'S REQUEST. USE THE \"invoke.sh\" LAUNCHER TO RUN LATER **\n")
+                print(
+                    '\n** CANCELLED AT USER\'S REQUEST. USE THE "invoke.sh" LAUNCHER TO RUN LATER **\n'
+                )
                 sys.exit(0)
 
         if opt.skip_support_models:

ldm/invoke/devices.py

@@ -1,19 +1,25 @@
+from __future__ import annotations
+
+from contextlib import nullcontext
+
 import torch
 from torch import autocast
-from contextlib import nullcontext
+
 from ldm.invoke.globals import Globals
 
-def choose_torch_device() -> str:
+CPU_DEVICE = torch.device("cpu")
+
+def choose_torch_device() -> torch.device:
     '''Convenience routine for guessing which GPU device to run model on'''
     if Globals.always_use_cpu:
-        return "cpu"
+        return CPU_DEVICE
     if torch.cuda.is_available():
-        return 'cuda'
+        return torch.device('cuda')
     if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
-        return 'mps'
-    return 'cpu'
+        return torch.device('mps')
+    return CPU_DEVICE
 
-def choose_precision(device) -> str:
+def choose_precision(device: torch.device) -> str:
     '''Returns an appropriate precision for the given torch device'''
     if device.type == 'cuda':
         device_name = torch.cuda.get_device_name(device)
@@ -21,7 +27,7 @@ def choose_precision(device) -> str:
             return 'float16'
     return 'float32'
 
-def torch_dtype(device) -> torch.dtype:
+def torch_dtype(device: torch.device) -> torch.dtype:
     if Globals.full_precision:
         return torch.float32
     if choose_precision(device) == 'float16':
@@ -36,3 +42,13 @@ def choose_autocast(precision):
     if precision == 'autocast' or precision == 'float16':
         return autocast
     return nullcontext
+
+def normalize_device(device: str | torch.device) -> torch.device:
+    """Ensure device has a device index defined, if appropriate."""
+    device = torch.device(device)
+    if device.index is None:
+        # cuda might be the only torch backend that currently uses the device index?
+        # I don't see anything like `current_device` for cpu or mps.
+        if device.type == 'cuda':
+            device = torch.device(device.type, torch.cuda.current_device())
+    return device

ldm/invoke/generator/diffusers_pipeline.py

@@ -29,6 +29,7 @@ from typing_extensions import ParamSpec
 from ldm.invoke.globals import Globals
 from ldm.models.diffusion.shared_invokeai_diffusion import InvokeAIDiffuserComponent, PostprocessingSettings
 from ldm.modules.textual_inversion_manager import TextualInversionManager
+from ..devices import normalize_device, CPU_DEVICE
 from ..offloading import LazilyLoadedModelGroup, FullyLoadedModelGroup, ModelGroup
 from ...models.diffusion.cross_attention_map_saving import AttentionMapSaver
 from ...modules.prompt_to_embeddings_converter import WeightedPromptFragmentsToEmbeddingsConverter
@@ -320,7 +321,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
                 if self.device.type == 'cpu' or self.device.type == 'mps':
                     mem_free = psutil.virtual_memory().free
                 elif self.device.type == 'cuda':
-                    mem_free, _ = torch.cuda.mem_get_info(self.device)
+                    mem_free, _ = torch.cuda.mem_get_info(normalize_device(self.device))
                 else:
                     raise ValueError(f"unrecognized device {self.device}")
                 # input tensor of [1, 4, h/8, w/8]
@@ -381,9 +382,10 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
         self._model_group.ready()
 
     def to(self, torch_device: Optional[Union[str, torch.device]] = None):
+        # overridden method; types match the superclass.
         if torch_device is None:
             return self
-        self._model_group.set_device(torch_device)
+        self._model_group.set_device(torch.device(torch_device))
         self._model_group.ready()
 
     @property
@@ -690,8 +692,8 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
             if device.type == 'mps':
                 # workaround for torch MPS bug that has been fixed in https://github.com/kulinseth/pytorch/pull/222
                 # TODO remove this workaround once kulinseth#222 is merged to pytorch mainline
-                self.vae.to('cpu')
-                init_image = init_image.to('cpu')
+                self.vae.to(CPU_DEVICE)
+                init_image = init_image.to(CPU_DEVICE)
             else:
                 self._model_group.load(self.vae)
             init_latent_dist = self.vae.encode(init_image).latent_dist

ldm/invoke/globals.py

@@ -54,6 +54,9 @@ Globals.full_precision = False
 # whether we should convert ckpt files into diffusers models on the fly
 Globals.ckpt_convert = False
 
+# logging tokenization everywhere
+Globals.log_tokenization = False
+
 def global_config_file()->Path:
     return Path(Globals.root, Globals.config_dir, Globals.models_file)
 

ldm/invoke/merge_diffusers.py

@@ -80,8 +80,8 @@ def merge_diffusion_models_and_commit(
     merged_model_name = name for new model
     alpha  - The interpolation parameter. Ranges from 0 to 1.  It affects the ratio in which the checkpoints are merged. A 0.8 alpha
                would mean that the first model checkpoints would affect the final result far less than an alpha of 0.2
-    interp - The interpolation method to use for the merging. Supports "sigmoid", "inv_sigmoid", "add_difference" and None.
-               Passing None uses the default interpolation which is weighted sum interpolation. For merging three checkpoints, only "add_difference" is supported.
+    interp - The interpolation method to use for the merging. Supports "weighted_average", "sigmoid", "inv_sigmoid", "add_difference" and None.
+               Passing None uses the default interpolation which is weighted sum interpolation. For merging three checkpoints, only "add_difference" is supported. Add_difference is A+(B-C).
     force  - Whether to ignore mismatch in model_config.json for the current models. Defaults to False.
 
     **kwargs - the default DiffusionPipeline.get_config_dict kwargs:
@@ -174,7 +174,7 @@ def _parse_args() -> Namespace:
 
 # ------------------------- GUI HERE -------------------------
 class mergeModelsForm(npyscreen.FormMultiPageAction):
-    interpolations = ["weighted_sum", "sigmoid", "inv_sigmoid", "add_difference"]
+    interpolations = ["weighted_sum", "sigmoid", "inv_sigmoid"]
 
     def __init__(self, parentApp, name):
         self.parentApp = parentApp
@@ -293,8 +293,8 @@ class mergeModelsForm(npyscreen.FormMultiPageAction):
         self.alpha = self.add_widget_intelligent(
             FloatTitleSlider,
             name="Weight (alpha) to assign to second and third models:",
-            out_of=1,
-            step=0.05,
+            out_of=1.0,
+            step=0.01,
             lowest=0,
             value=0.5,
             labelColor="CONTROL",
@@ -311,7 +311,7 @@ class mergeModelsForm(npyscreen.FormMultiPageAction):
         self.merged_model_name.value = merged_model_name
 
         if selected_model3 > 0:
-            self.merge_method.values = (["add_difference"],)
+            self.merge_method.values = ['add_difference ( A+(B-C) )']
             self.merged_model_name.value += f"+{models[selected_model3]}"
         else:
             self.merge_method.values = self.interpolations
@@ -337,11 +337,14 @@ class mergeModelsForm(npyscreen.FormMultiPageAction):
         ]
         if self.model3.value[0] > 0:
             models.append(model_names[self.model3.value[0] - 1])
+            interp='add_difference'
+        else:
+            interp=self.interpolations[self.merge_method.value[0]]
 
         args = dict(
             models=models,
             alpha=self.alpha.value,
-            interp=self.interpolations[self.merge_method.value[0]],
+            interp=interp,
             force=self.force.value,
             merged_model_name=self.merged_model_name.value,
         )

ldm/invoke/model_manager.py

@@ -32,6 +32,7 @@ from omegaconf import OmegaConf
 from omegaconf.dictconfig import DictConfig
 from picklescan.scanner import scan_file_path
 
+from ldm.invoke.devices import CPU_DEVICE
 from ldm.invoke.generator.diffusers_pipeline import \
     StableDiffusionGeneratorPipeline
 from ldm.invoke.globals import (Globals, global_cache_dir)
@@ -54,7 +55,7 @@ class ModelManager(object):
     def __init__(
         self,
         config: OmegaConf,
-        device_type: str | torch.device = "cpu",
+        device_type: torch.device = CPU_DEVICE,
         precision: str = "float16",
         max_loaded_models=DEFAULT_MAX_MODELS,
         sequential_offload = False
@@ -672,7 +673,7 @@ class ModelManager(object):
         """
         if str(weights).startswith(("http:", "https:")):
             model_name = model_name or url_attachment_name(weights)
-            
+
         weights_path = self._resolve_path(weights, "models/ldm/stable-diffusion-v1")
         config_path  = self._resolve_path(config, "configs/stable-diffusion")
 
@@ -1112,25 +1113,25 @@ class ModelManager(object):
         self.models.pop(model_name, None)
 
     def _model_to_cpu(self, model):
-        if self.device == "cpu":
+        if self.device == CPU_DEVICE:
             return model
 
         if isinstance(model, StableDiffusionGeneratorPipeline):
             model.offload_all()
             return model
 
-        model.cond_stage_model.device = "cpu"
-        model.to("cpu")
+        model.cond_stage_model.device = CPU_DEVICE
+        model.to(CPU_DEVICE)
 
         for submodel in ("first_stage_model", "cond_stage_model", "model"):
             try:
-                getattr(model, submodel).to("cpu")
+                getattr(model, submodel).to(CPU_DEVICE)
             except AttributeError:
                 pass
         return model
 
     def _model_from_cpu(self, model):
-        if self.device == "cpu":
+        if self.device == CPU_DEVICE:
             return model
 
         if isinstance(model, StableDiffusionGeneratorPipeline):

ldm/invoke/readline.py

@@ -60,7 +60,7 @@ COMMANDS = (
     '--text_mask','-tm',
     '!fix','!fetch','!replay','!history','!search','!clear',
     '!models','!switch','!import_model','!optimize_model','!convert_model','!edit_model','!del_model',
-    '!mask',
+    '!mask','!triggers',
     )
 MODEL_COMMANDS = (
     '!switch',

ldm/modules/textual_inversion_manager.py

@@ -1,11 +1,12 @@
 import os
 import traceback
-from typing import Optional
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Union
 
 import torch
-from dataclasses import dataclass
 from picklescan.scanner import scan_file_path
-from transformers import CLIPTokenizer, CLIPTextModel
+from transformers import CLIPTextModel, CLIPTokenizer
 
 from ldm.invoke.concepts_lib import HuggingFaceConceptsLibrary
 
@@ -21,11 +22,14 @@ class TextualInversion:
     def embedding_vector_length(self) -> int:
         return self.embedding.shape[0]
 
-class TextualInversionManager():
-    def __init__(self,
-                 tokenizer: CLIPTokenizer,
-                 text_encoder: CLIPTextModel,
-                 full_precision: bool=True):
+
+class TextualInversionManager:
+    def __init__(
+        self,
+        tokenizer: CLIPTokenizer,
+        text_encoder: CLIPTextModel,
+        full_precision: bool = True,
+    ):
         self.tokenizer = tokenizer
         self.text_encoder = text_encoder
         self.full_precision = full_precision
@@ -38,47 +42,73 @@ class TextualInversionManager():
             if concept_name in self.hf_concepts_library.concepts_loaded:
                 continue
             trigger = self.hf_concepts_library.concept_to_trigger(concept_name)
-            if self.has_textual_inversion_for_trigger_string(trigger) \
-               or self.has_textual_inversion_for_trigger_string(concept_name) \
-               or self.has_textual_inversion_for_trigger_string(f'<{concept_name}>'):  # in case a token with literal angle brackets encountered
-                print(f'>> Loaded local embedding for trigger {concept_name}')
+            if (
+                self.has_textual_inversion_for_trigger_string(trigger)
+                or self.has_textual_inversion_for_trigger_string(concept_name)
+                or self.has_textual_inversion_for_trigger_string(f"<{concept_name}>")
+            ):  # in case a token with literal angle brackets encountered
+                print(f">> Loaded local embedding for trigger {concept_name}")
                 continue
             bin_file = self.hf_concepts_library.get_concept_model_path(concept_name)
             if not bin_file:
                 continue
-            print(f'>> Loaded remote embedding for trigger {concept_name}')
+            print(f">> Loaded remote embedding for trigger {concept_name}")
             self.load_textual_inversion(bin_file)
-            self.hf_concepts_library.concepts_loaded[concept_name]=True
+            self.hf_concepts_library.concepts_loaded[concept_name] = True
 
     def get_all_trigger_strings(self) -> list[str]:
         return [ti.trigger_string for ti in self.textual_inversions]
 
-    def load_textual_inversion(self, ckpt_path, defer_injecting_tokens: bool=False):
-        if str(ckpt_path).endswith('.DS_Store'):
+    def load_textual_inversion(self, ckpt_path: Union[str,Path], defer_injecting_tokens: bool = False):
+        ckpt_path = Path(ckpt_path)
+        if str(ckpt_path).endswith(".DS_Store"):
             return
         try:
-            scan_result = scan_file_path(ckpt_path)
+            scan_result = scan_file_path(str(ckpt_path))
             if scan_result.infected_files == 1:
-                print(f'\n### Security Issues Found in Model: {scan_result.issues_count}')
-                print('### For your safety, InvokeAI will not load this embed.')
+                print(
+                    f"\n### Security Issues Found in Model: {scan_result.issues_count}"
+                )
+                print("### For your safety, InvokeAI will not load this embed.")
                 return
         except Exception:
-            print(f"### WARNING::: Invalid or corrupt embeddings found. Ignoring: {ckpt_path}")
+            print(
+                f"### {ckpt_path.parents[0].name}/{ckpt_path.name} is damaged or corrupt."
+            )
+            return
+
+        embedding_info = self._parse_embedding(str(ckpt_path))
+
+        if embedding_info is None:
+            # We've already put out an error message about the bad embedding in _parse_embedding, so just return.
+            return
+        elif (
+            self.text_encoder.get_input_embeddings().weight.data[0].shape[0]
+            != embedding_info["embedding"].shape[0]
+        ):
+            print(
+                f"** Notice: {ckpt_path.parents[0].name}/{ckpt_path.name} was trained on a model with a different token dimension. It can't be used with this model."
+            )
             return
 
-        embedding_info = self._parse_embedding(ckpt_path)
         if embedding_info:
             try:
-                self._add_textual_inversion(embedding_info['name'],
-                                            embedding_info['embedding'],
-                                            defer_injecting_tokens=defer_injecting_tokens)
+                self._add_textual_inversion(
+                    embedding_info["name"],
+                    embedding_info["embedding"],
+                    defer_injecting_tokens=defer_injecting_tokens,
+                )
             except ValueError as e:
                 print(f'   | Ignoring incompatible embedding {embedding_info["name"]}')
-                print(f'   | The error was {str(e)}')
+                print(f"   | The error was {str(e)}")
         else:
-            print(f'>> Failed to load embedding located at {ckpt_path}. Unsupported file.')
+            print(
+                f">> Failed to load embedding located at {str(ckpt_path)}. Unsupported file."
+            )
 
-    def _add_textual_inversion(self, trigger_str, embedding, defer_injecting_tokens=False) -> TextualInversion:
+    def _add_textual_inversion(
+        self, trigger_str, embedding, defer_injecting_tokens=False
+    ) -> TextualInversion:
         """
         Add a textual inversion to be recognised.
         :param trigger_str: The trigger text in the prompt that activates this textual inversion. If unknown to the embedder's tokenizer, will be added.
@@ -86,46 +116,59 @@ class TextualInversionManager():
         :return: The token id for the added embedding, either existing or newly-added.
         """
         if trigger_str in [ti.trigger_string for ti in self.textual_inversions]:
-            print(f">> TextualInversionManager refusing to overwrite already-loaded token '{trigger_str}'")
+            print(
+                f">> TextualInversionManager refusing to overwrite already-loaded token '{trigger_str}'"
+            )
             return
         if not self.full_precision:
             embedding = embedding.half()
         if len(embedding.shape) == 1:
             embedding = embedding.unsqueeze(0)
         elif len(embedding.shape) > 2:
-            raise ValueError(f"TextualInversionManager cannot add {trigger_str} because the embedding shape {embedding.shape} is incorrect. The embedding must have shape [token_dim] or [V, token_dim] where V is vector length and token_dim is 768 for SD1 or 1280 for SD2.")
+            raise ValueError(
+                f"TextualInversionManager cannot add {trigger_str} because the embedding shape {embedding.shape} is incorrect. The embedding must have shape [token_dim] or [V, token_dim] where V is vector length and token_dim is 768 for SD1 or 1280 for SD2."
+            )
 
         try:
-            ti = TextualInversion(
-                trigger_string=trigger_str,
-                embedding=embedding
-            )
+            ti = TextualInversion(trigger_string=trigger_str, embedding=embedding)
             if not defer_injecting_tokens:
                 self._inject_tokens_and_assign_embeddings(ti)
             self.textual_inversions.append(ti)
             return ti
 
         except ValueError as e:
-            if str(e).startswith('Warning'):
+            if str(e).startswith("Warning"):
                 print(f">> {str(e)}")
             else:
                 traceback.print_exc()
-                print(f">> TextualInversionManager was unable to add a textual inversion with trigger string {trigger_str}.")
+                print(
+                    f">> TextualInversionManager was unable to add a textual inversion with trigger string {trigger_str}."
+                )
                 raise
 
     def _inject_tokens_and_assign_embeddings(self, ti: TextualInversion) -> int:
-
         if ti.trigger_token_id is not None:
-            raise ValueError(f"Tokens already injected for textual inversion with trigger '{ti.trigger_string}'")
+            raise ValueError(
+                f"Tokens already injected for textual inversion with trigger '{ti.trigger_string}'"
+            )
 
-        trigger_token_id = self._get_or_create_token_id_and_assign_embedding(ti.trigger_string, ti.embedding[0])
+        trigger_token_id = self._get_or_create_token_id_and_assign_embedding(
+            ti.trigger_string, ti.embedding[0]
+        )
 
         if ti.embedding_vector_length > 1:
             # for embeddings with vector length > 1
-            pad_token_strings = [ti.trigger_string + "-!pad-" + str(pad_index) for pad_index in range(1, ti.embedding_vector_length)]
+            pad_token_strings = [
+                ti.trigger_string + "-!pad-" + str(pad_index)
+                for pad_index in range(1, ti.embedding_vector_length)
+            ]
             # todo: batched UI for faster loading when vector length >2
-            pad_token_ids = [self._get_or_create_token_id_and_assign_embedding(pad_token_str, ti.embedding[1 + i]) \
-                                 for (i, pad_token_str) in enumerate(pad_token_strings)]
+            pad_token_ids = [
+                self._get_or_create_token_id_and_assign_embedding(
+                    pad_token_str, ti.embedding[1 + i]
+                )
+                for (i, pad_token_str) in enumerate(pad_token_strings)
+            ]
         else:
             pad_token_ids = []
 
@@ -133,7 +176,6 @@ class TextualInversionManager():
         ti.pad_token_ids = pad_token_ids
         return ti.trigger_token_id
 
-
     def has_textual_inversion_for_trigger_string(self, trigger_string: str) -> bool:
         try:
             ti = self.get_textual_inversion_for_trigger_string(trigger_string)
@@ -141,32 +183,43 @@ class TextualInversionManager():
         except StopIteration:
             return False
 
-
-    def get_textual_inversion_for_trigger_string(self, trigger_string: str) -> TextualInversion:
-        return next(ti for ti in self.textual_inversions if ti.trigger_string == trigger_string)
-
+    def get_textual_inversion_for_trigger_string(
+        self, trigger_string: str
+    ) -> TextualInversion:
+        return next(
+            ti for ti in self.textual_inversions if ti.trigger_string == trigger_string
+        )
 
     def get_textual_inversion_for_token_id(self, token_id: int) -> TextualInversion:
-        return next(ti for ti in self.textual_inversions if ti.trigger_token_id == token_id)
+        return next(
+            ti for ti in self.textual_inversions if ti.trigger_token_id == token_id
+        )
 
-    def create_deferred_token_ids_for_any_trigger_terms(self, prompt_string: str) -> list[int]:
+    def create_deferred_token_ids_for_any_trigger_terms(
+        self, prompt_string: str
+    ) -> list[int]:
         injected_token_ids = []
         for ti in self.textual_inversions:
             if ti.trigger_token_id is None and ti.trigger_string in prompt_string:
                 if ti.embedding_vector_length > 1:
-                    print(f">> Preparing tokens for textual inversion {ti.trigger_string}...")
+                    print(
+                        f">> Preparing tokens for textual inversion {ti.trigger_string}..."
+                    )
                 try:
                     self._inject_tokens_and_assign_embeddings(ti)
                 except ValueError as e:
-                    print(f'   | Ignoring incompatible embedding trigger {ti.trigger_string}')
-                    print(f'   | The error was {str(e)}')
+                    print(
+                        f"   | Ignoring incompatible embedding trigger {ti.trigger_string}"
+                    )
+                    print(f"   | The error was {str(e)}")
                     continue
                 injected_token_ids.append(ti.trigger_token_id)
                 injected_token_ids.extend(ti.pad_token_ids)
         return injected_token_ids
 
-
-    def expand_textual_inversion_token_ids_if_necessary(self, prompt_token_ids: list[int]) -> list[int]:
+    def expand_textual_inversion_token_ids_if_necessary(
+        self, prompt_token_ids: list[int]
+    ) -> list[int]:
         """
         Insert padding tokens as necessary into the passed-in list of token ids to match any textual inversions it includes.
 
@@ -181,20 +234,31 @@ class TextualInversionManager():
             raise ValueError("prompt_token_ids must not start with bos_token_id")
         if prompt_token_ids[-1] == self.tokenizer.eos_token_id:
             raise ValueError("prompt_token_ids must not end with eos_token_id")
-        textual_inversion_trigger_token_ids = [ti.trigger_token_id for ti in self.textual_inversions]
+        textual_inversion_trigger_token_ids = [
+            ti.trigger_token_id for ti in self.textual_inversions
+        ]
         prompt_token_ids = prompt_token_ids.copy()
         for i, token_id in reversed(list(enumerate(prompt_token_ids))):
             if token_id in textual_inversion_trigger_token_ids:
-                textual_inversion = next(ti for ti in self.textual_inversions if ti.trigger_token_id == token_id)
-                for pad_idx in range(0, textual_inversion.embedding_vector_length-1):
-                    prompt_token_ids.insert(i+pad_idx+1, textual_inversion.pad_token_ids[pad_idx])
+                textual_inversion = next(
+                    ti
+                    for ti in self.textual_inversions
+                    if ti.trigger_token_id == token_id
+                )
+                for pad_idx in range(0, textual_inversion.embedding_vector_length - 1):
+                    prompt_token_ids.insert(
+                        i + pad_idx + 1, textual_inversion.pad_token_ids[pad_idx]
+                    )
 
         return prompt_token_ids
 
-
-    def _get_or_create_token_id_and_assign_embedding(self, token_str: str, embedding: torch.Tensor) -> int:
+    def _get_or_create_token_id_and_assign_embedding(
+        self, token_str: str, embedding: torch.Tensor
+    ) -> int:
         if len(embedding.shape) != 1:
-            raise ValueError("Embedding has incorrect shape - must be [token_dim] where token_dim is 768 for SD1 or 1280 for SD2")
+            raise ValueError(
+                "Embedding has incorrect shape - must be [token_dim] where token_dim is 768 for SD1 or 1280 for SD2"
+            )
         existing_token_id = self.tokenizer.convert_tokens_to_ids(token_str)
         if existing_token_id == self.tokenizer.unk_token_id:
             num_tokens_added = self.tokenizer.add_tokens(token_str)
@@ -207,66 +271,79 @@ class TextualInversionManager():
         token_id = self.tokenizer.convert_tokens_to_ids(token_str)
         if token_id == self.tokenizer.unk_token_id:
             raise RuntimeError(f"Unable to find token id for token '{token_str}'")
-        if self.text_encoder.get_input_embeddings().weight.data[token_id].shape != embedding.shape:
-            raise ValueError(f"Warning. Cannot load embedding for {token_str}. It was trained on a model with token dimension {embedding.shape[0]}, but the current model has token dimension {self.text_encoder.get_input_embeddings().weight.data[token_id].shape[0]}.")
+        if (
+            self.text_encoder.get_input_embeddings().weight.data[token_id].shape
+            != embedding.shape
+        ):
+            raise ValueError(
+                f"Warning. Cannot load embedding for {token_str}. It was trained on a model with token dimension {embedding.shape[0]}, but the current model has token dimension {self.text_encoder.get_input_embeddings().weight.data[token_id].shape[0]}."
+            )
         self.text_encoder.get_input_embeddings().weight.data[token_id] = embedding
 
         return token_id
 
     def _parse_embedding(self, embedding_file: str):
-        file_type = embedding_file.split('.')[-1]
-        if file_type == 'pt':
+        file_type = embedding_file.split(".")[-1]
+        if file_type == "pt":
             return self._parse_embedding_pt(embedding_file)
-        elif file_type == 'bin':
+        elif file_type == "bin":
             return self._parse_embedding_bin(embedding_file)
         else:
-            print(f'>> Not a recognized embedding file: {embedding_file}')
+            print(f">> Not a recognized embedding file: {embedding_file}")
+            return None
 
     def _parse_embedding_pt(self, embedding_file):
-        embedding_ckpt = torch.load(embedding_file, map_location='cpu')
+        embedding_ckpt = torch.load(embedding_file, map_location="cpu")
         embedding_info = {}
 
         # Check if valid embedding file
-        if 'string_to_token' and 'string_to_param' in embedding_ckpt:
-
+        if "string_to_token" and "string_to_param" in embedding_ckpt:
             # Catch variants that do not have the expected keys or values.
             try:
-                embedding_info['name'] = embedding_ckpt['name'] or os.path.basename(os.path.splitext(embedding_file)[0])
+                embedding_info["name"] = embedding_ckpt["name"] or os.path.basename(
+                    os.path.splitext(embedding_file)[0]
+                )
 
                 # Check num of embeddings and warn user only the first will be used
-                embedding_info['num_of_embeddings'] = len(embedding_ckpt["string_to_token"])
-                if embedding_info['num_of_embeddings'] > 1:
-                    print('>> More than 1 embedding found. Will use the first one')
+                embedding_info["num_of_embeddings"] = len(
+                    embedding_ckpt["string_to_token"]
+                )
+                if embedding_info["num_of_embeddings"] > 1:
+                    print(">> More than 1 embedding found. Will use the first one")
 
-                embedding = list(embedding_ckpt['string_to_param'].values())[0]
-            except (AttributeError,KeyError):
+                embedding = list(embedding_ckpt["string_to_param"].values())[0]
+            except (AttributeError, KeyError):
                 return self._handle_broken_pt_variants(embedding_ckpt, embedding_file)
 
-            embedding_info['embedding'] = embedding
-            embedding_info['num_vectors_per_token'] = embedding.size()[0]
-            embedding_info['token_dim'] = embedding.size()[1]
+            embedding_info["embedding"] = embedding
+            embedding_info["num_vectors_per_token"] = embedding.size()[0]
+            embedding_info["token_dim"] = embedding.size()[1]
 
             try:
-                embedding_info['trained_steps'] = embedding_ckpt['step']
-                embedding_info['trained_model_name'] = embedding_ckpt['sd_checkpoint_name']
-                embedding_info['trained_model_checksum'] = embedding_ckpt['sd_checkpoint']
+                embedding_info["trained_steps"] = embedding_ckpt["step"]
+                embedding_info["trained_model_name"] = embedding_ckpt[
+                    "sd_checkpoint_name"
+                ]
+                embedding_info["trained_model_checksum"] = embedding_ckpt[
+                    "sd_checkpoint"
+                ]
             except AttributeError:
                 print(">> No Training Details Found. Passing ...")
 
         # .pt files found at https://cyberes.github.io/stable-diffusion-textual-inversion-models/
         # They are actually .bin files
-        elif len(embedding_ckpt.keys())==1:
-            print('>> Detected .bin file masquerading as .pt file')
+        elif len(embedding_ckpt.keys()) == 1:
+            print(">> Detected .bin file masquerading as .pt file")
             embedding_info = self._parse_embedding_bin(embedding_file)
 
         else:
-            print('>> Invalid embedding format')
+            print(">> Invalid embedding format")
             embedding_info = None
 
         return embedding_info
 
     def _parse_embedding_bin(self, embedding_file):
-        embedding_ckpt = torch.load(embedding_file, map_location='cpu')
+        embedding_ckpt = torch.load(embedding_file, map_location="cpu")
         embedding_info = {}
 
         if list(embedding_ckpt.keys()) == 0:
@@ -274,27 +351,45 @@ class TextualInversionManager():
             embedding_info = None
         else:
             for token in list(embedding_ckpt.keys()):
-                embedding_info['name'] = token or os.path.basename(os.path.splitext(embedding_file)[0])
-                embedding_info['embedding'] = embedding_ckpt[token]
-                embedding_info['num_vectors_per_token'] = 1 # All Concepts seem to default to 1
-                embedding_info['token_dim'] = embedding_info['embedding'].size()[0]
+                embedding_info["name"] = token or os.path.basename(
+                    os.path.splitext(embedding_file)[0]
+                )
+                embedding_info["embedding"] = embedding_ckpt[token]
+                embedding_info[
+                    "num_vectors_per_token"
+                ] = 1  # All Concepts seem to default to 1
+                embedding_info["token_dim"] = embedding_info["embedding"].size()[0]
 
         return embedding_info
 
-    def _handle_broken_pt_variants(self, embedding_ckpt:dict, embedding_file:str)->dict:
-        '''
+    def _handle_broken_pt_variants(
+        self, embedding_ckpt: dict, embedding_file: str
+    ) -> dict:
+        """
         This handles the broken .pt file variants. We only know of one at present.
-        '''
+        """
         embedding_info = {}
-        if isinstance(list(embedding_ckpt['string_to_token'].values())[0],torch.Tensor):
-            print('>> Detected .pt file variant 1') # example at https://github.com/invoke-ai/InvokeAI/issues/1829
-            for token in list(embedding_ckpt['string_to_token'].keys()):
-                embedding_info['name'] = token if token != '*' else os.path.basename(os.path.splitext(embedding_file)[0])
-                embedding_info['embedding'] = embedding_ckpt['string_to_param'].state_dict()[token]
-                embedding_info['num_vectors_per_token'] = embedding_info['embedding'].shape[0]
-                embedding_info['token_dim'] = embedding_info['embedding'].size()[0]
+        if isinstance(
+            list(embedding_ckpt["string_to_token"].values())[0], torch.Tensor
+        ):
+            print(
+                ">> Detected .pt file variant 1"
+            )  # example at https://github.com/invoke-ai/InvokeAI/issues/1829
+            for token in list(embedding_ckpt["string_to_token"].keys()):
+                embedding_info["name"] = (
+                    token
+                    if token != "*"
+                    else os.path.basename(os.path.splitext(embedding_file)[0])
+                )
+                embedding_info["embedding"] = embedding_ckpt[
+                    "string_to_param"
+                ].state_dict()[token]
+                embedding_info["num_vectors_per_token"] = embedding_info[
+                    "embedding"
+                ].shape[0]
+                embedding_info["token_dim"] = embedding_info["embedding"].size()[0]
         else:
-            print('>> Invalid embedding format')
+            print(">> Invalid embedding format")
             embedding_info = None
 
         return embedding_info