Merge branch 'main' into refactor/model-manager-2

invokeai/backend/install/model_install_backend.py

@@ -254,7 +254,13 @@ class ModelInstall(object):
         elif path.is_dir() and any(
             [
                 (path / x).exists()
-                for x in {"config.json", "model_index.json", "learned_embeds.bin", "pytorch_lora_weights.bin"}
+                for x in {
+                    "config.json",
+                    "model_index.json",
+                    "learned_embeds.bin",
+                    "pytorch_lora_weights.bin",
+                    "pytorch_lora_weights.safetensors",
+                }
             ]
         ):
             models_installed.update({str(model_path_id_or_url): self._install_path(path)})
@@ -357,7 +363,7 @@ class ModelInstall(object):
                 for suffix in ["safetensors", "bin"]:
                     if f"{prefix}pytorch_lora_weights.{suffix}" in files:
                         location = self._download_hf_model(
-                            repo_id, ["pytorch_lora_weights.bin"], staging, subfolder=subfolder
+                            repo_id, [f"pytorch_lora_weights.{suffix}"], staging, subfolder=subfolder
                         )  # LoRA
                         break
                     elif (

invokeai/backend/model_management/lora.py

@@ -166,6 +166,15 @@ class ModelPatcher:
         init_tokens_count = None
         new_tokens_added = None
 
+        # TODO: This is required since Transformers 4.32 see
+        # https://github.com/huggingface/transformers/pull/25088
+        # More information by NVIDIA:
+        # https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#requirements-tc
+        # This value might need to be changed in the future and take the GPUs model into account as there seem
+        # to be ideal values for different GPUS. This value is temporary!
+        # For references to the current discussion please see https://github.com/invoke-ai/InvokeAI/pull/4817
+        pad_to_multiple_of = 8
+
         try:
             # HACK: The CLIPTokenizer API does not include a way to remove tokens after calling add_tokens(...). As a
             # workaround, we create a full copy of `tokenizer` so that its original behavior can be restored after
@@ -175,7 +184,7 @@ class ModelPatcher:
             # but a pickle roundtrip was found to be much faster (1 sec vs. 0.05 secs).
             ti_tokenizer = pickle.loads(pickle.dumps(tokenizer))
             ti_manager = TextualInversionManager(ti_tokenizer)
-            init_tokens_count = text_encoder.resize_token_embeddings(None).num_embeddings
+            init_tokens_count = text_encoder.resize_token_embeddings(None, pad_to_multiple_of).num_embeddings
 
             def _get_trigger(ti_name, index):
                 trigger = ti_name
@@ -190,7 +199,7 @@ class ModelPatcher:
                     new_tokens_added += ti_tokenizer.add_tokens(_get_trigger(ti_name, i))
 
             # modify text_encoder
-            text_encoder.resize_token_embeddings(init_tokens_count + new_tokens_added)
+            text_encoder.resize_token_embeddings(init_tokens_count + new_tokens_added, pad_to_multiple_of)
             model_embeddings = text_encoder.get_input_embeddings()
 
             for ti_name, ti in ti_list:
@@ -222,7 +231,7 @@ class ModelPatcher:
 
         finally:
             if init_tokens_count and new_tokens_added:
-                text_encoder.resize_token_embeddings(init_tokens_count)
+                text_encoder.resize_token_embeddings(init_tokens_count, pad_to_multiple_of)
 
     @classmethod
     @contextmanager

invokeai/backend/model_management/model_probe.py

@@ -183,12 +183,13 @@ class ModelProbe(object):
         if model:
             class_name = model.__class__.__name__
         else:
+            for suffix in ["bin", "safetensors"]:
+                if (folder_path / f"learned_embeds.{suffix}").exists():
+                    return ModelType.TextualInversion
+                if (folder_path / f"pytorch_lora_weights.{suffix}").exists():
+                    return ModelType.Lora
             if (folder_path / "unet/model.onnx").exists():
                 return ModelType.ONNX
-            if (folder_path / "learned_embeds.bin").exists():
-                return ModelType.TextualInversion
-            if (folder_path / "pytorch_lora_weights.bin").exists():
-                return ModelType.Lora
             if (folder_path / "image_encoder.txt").exists():
                 return ModelType.IPAdapter
 

invokeai/backend/model_management/models/lora.py

@@ -68,8 +68,9 @@ class LoRAModel(ModelBase):
             raise ModelNotFoundException()
 
         if os.path.isdir(path):
-            if os.path.exists(os.path.join(path, "pytorch_lora_weights.bin")):
-                return LoRAModelFormat.Diffusers
+            for ext in ["safetensors", "bin"]:
+                if os.path.exists(os.path.join(path, f"pytorch_lora_weights.{ext}")):
+                    return LoRAModelFormat.Diffusers
 
         if os.path.isfile(path):
             if any([path.endswith(f".{ext}") for ext in ["safetensors", "ckpt", "pt"]]):
@@ -86,8 +87,10 @@ class LoRAModel(ModelBase):
         base_model: BaseModelType,
     ) -> str:
         if cls.detect_format(model_path) == LoRAModelFormat.Diffusers:
-            # TODO: add diffusers lora when it stabilizes a bit
-            raise NotImplementedError("Diffusers lora not supported")
+            for ext in ["safetensors", "bin"]:  # return path to the safetensors file inside the folder
+                path = Path(model_path, f"pytorch_lora_weights.{ext}")
+                if path.exists():
+                    return path
         else:
             return model_path