From 8702a63197de17e122799c8db4b5a39c043d19ab Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lstein@gmail.com>
Date: Fri, 10 Nov 2023 17:38:34 -0500
Subject: [PATCH] add support for downloading and installing LCM lora diffusers
 models

---
 invokeai/backend/install/model_install_backend.py | 10 ++++++++--
 invokeai/backend/model_management/model_probe.py  |  9 +++++----
 invokeai/backend/model_management/models/lora.py  | 11 +++++++----
 3 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/invokeai/backend/install/model_install_backend.py b/invokeai/backend/install/model_install_backend.py
index 9784aa0ac2..34526cfaf3 100644
--- a/invokeai/backend/install/model_install_backend.py
+++ b/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 (
diff --git a/invokeai/backend/model_management/model_probe.py b/invokeai/backend/model_management/model_probe.py
index 19d64b035f..e82be8d069 100644
--- a/invokeai/backend/model_management/model_probe.py
+++ b/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
 
diff --git a/invokeai/backend/model_management/models/lora.py b/invokeai/backend/model_management/models/lora.py
index 6ececa7df1..43a24275d1 100644
--- a/invokeai/backend/model_management/models/lora.py
+++ b/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