[Enhancement] distinguish v1 from v2 LoRA models (#3175)

# Distinguish LoRA/LyCORIS files based on what version of SD they were
built on top of

- Attempting to run a prompt with a LoRA based on SD v1.X against a
model based on v2.X will now throw an `IncompatibleModelException`. To
import this exception:
`from ldm.modules.lora_manager import IncompatibleModelException` (maybe
this should be defined in ModelManager?)
    
- Enhance `LoraManager.list_loras()` to accept an optional integer
argument, `token_vector_length`. This will filter the returned LoRA
models to return only those that match the indicated length. Use:
      ```
      768 => for models based on SD v1.X
      1024 => for models based on SD v2.X
      ```
Note that this filtering requires each LoRA file to be opened by
`torch.safetensors`. It will take ~8s to scan a directory of 40 files.
    
- Added new static methods to `ldm.modules.kohya_lora_manager`:
      - check_model_compatibility()
      - vector_length_from_checkpoint()
      - vector_length_from_checkpoint_file()

- You can now create subdirectories within the `loras` directory and
organize the model files.

Author: lstein

invokeai/backend/invoke_ai_web_server.py

@@ -37,11 +37,11 @@
     Globals,
     global_converted_ckpts_dir,
     global_models_dir,
-    global_lora_models_dir,
 )
 from ldm.invoke.pngwriter import PngWriter, retrieve_metadata
 from compel.prompt_parser import Blend
 from ldm.invoke.merge_diffusers import merge_diffusion_models
+from ldm.modules.lora_manager import LoraManager
 
 # Loading Arguments
 opt = Args()
@@ -523,20 +523,12 @@ def merge_diffusers_models(model_merge_info: dict):
         @socketio.on("getLoraModels")
         def get_lora_models():
             try:
-                lora_path = global_lora_models_dir()
-                loras = []
-                for root, _, files in os.walk(lora_path):
-                    models = [
-                        Path(root, x)
-                        for x in files
-                        if Path(x).suffix in [".ckpt", ".pt", ".safetensors"]
-                    ]
-                    loras = loras + models
-
+                model = self.generate.model
+                lora_mgr = LoraManager(model)
+                loras = lora_mgr.list_compatible_loras()
                 found_loras = []
-                for lora in sorted(loras, key=lambda s: s.stem.lower()):
-                    location = str(lora.resolve()).replace("\\", "/")
-                    found_loras.append({"name": lora.stem, "location": location})
+                for lora in sorted(loras, key=str.casefold):
+                    found_loras.append({"name":lora,"location":str(loras[lora])})
                 socketio.emit("foundLoras", found_loras)
             except Exception as e:
                 self.handle_exceptions(e)

invokeai/configs/INITIAL_MODELS.yaml

@@ -80,7 +80,8 @@ trinart-2.0:
      repo_id: stabilityai/sd-vae-ft-mse
    recommended: False
 waifu-diffusion-1.4:
-   description: An SD-1.5 model trained on 680k anime/manga-style images (2.13 GB)
+   description: An SD-2.1 model trained on 5.4M anime/manga-style images (4.27 GB)
+   revision: main
    repo_id: hakurei/waifu-diffusion
    format: diffusers
    vae:

invokeai/frontend/dist/assets/index-f56b39bc.js renamed to invokeai/frontend/dist/assets/index-b12e648e.js

@@ -5,7 +5,7 @@
     <meta name="viewport" content="width=device-width, initial-scale=1.0" />
     <title>InvokeAI - A Stable Diffusion Toolkit</title>
     <link rel="shortcut icon" type="icon" href="./assets/favicon-0d253ced.ico" />
-    <script type="module" crossorigin src="./assets/index-f56b39bc.js"></script>
+    <script type="module" crossorigin src="./assets/index-b12e648e.js"></script>
     <link rel="stylesheet" href="./assets/index-2ab0eb58.css">
   </head>
 

invokeai/frontend/dist/index.html

@@ -33,6 +33,10 @@ import {
   setIntermediateImage,
 } from 'features/gallery/store/gallerySlice';
 
+import {
+  getLoraModels,
+  getTextualInversionTriggers,
+} from 'app/socketio/actions';
 import type { RootState } from 'app/store';
 import { addImageToStagingArea } from 'features/canvas/store/canvasSlice';
 import {
@@ -463,6 +467,8 @@ const makeSocketIOListeners = (
       const { model_name, model_list } = data;
       dispatch(setModelList(model_list));
       dispatch(setCurrentStatus(i18n.t('common.statusModelChanged')));
+      dispatch(getLoraModels());
+      dispatch(getTextualInversionTriggers());
       dispatch(setIsProcessing(false));
       dispatch(setIsCancelable(true));
       dispatch(

invokeai/frontend/src/app/socketio/listeners.ts

@@ -633,9 +633,8 @@ def process_image(image,seed):
         except RuntimeError:
             # Clear the CUDA cache on an exception
             self.clear_cuda_cache()
-
-            print(traceback.format_exc(), file=sys.stderr)
-            print(">> Could not generate image.")
+            print("** Could not generate image.")
+            raise
 
         toc = time.time()
         print("\n>> Usage stats:")

invokeai/frontend/stats.html

@@ -11,6 +11,7 @@
 
 import requests
 from diffusers import AutoencoderKL
+from diffusers import logging as dlogging
 from huggingface_hub import hf_hub_url
 from omegaconf import OmegaConf
 from omegaconf.dictconfig import DictConfig
@@ -295,13 +296,21 @@ def _download_diffusion_weights(
     mconfig: DictConfig, access_token: str, precision: str = "float32"
 ):
     repo_id = mconfig["repo_id"]
+    revision = mconfig.get('revision',None)
     model_class = (
         StableDiffusionGeneratorPipeline
         if mconfig.get("format", None) == "diffusers"
         else AutoencoderKL
     )
-    extra_arg_list = [{"revision": "fp16"}, {}] if precision == "float16" else [{}]
+    extra_arg_list = [{"revision": revision}] if revision \
+        else [{"revision": "fp16"}, {}] if precision == "float16" \
+             else [{}]
     path = None
+
+    # quench safety checker warnings
+    verbosity = dlogging.get_verbosity()
+    dlogging.set_verbosity_error()
+    
     for extra_args in extra_arg_list:
         try:
             path = download_from_hf(
@@ -317,6 +326,7 @@ def _download_diffusion_weights(
                 print(f"An unexpected error occurred while downloading the model: {e})")
         if path:
             break
+    dlogging.set_verbosity(verbosity)
     return path
 
 
@@ -447,6 +457,8 @@ def new_config_file_contents(
         stanza["description"] = mod["description"]
         stanza["repo_id"] = mod["repo_id"]
         stanza["format"] = mod["format"]
+        if "revision" in mod:
+            stanza["revision"] = mod["revision"]
         # diffusers don't need width and height (probably .ckpt doesn't either)
         # so we no longer require these in INITIAL_MODELS.yaml
         if "width" in mod:
@@ -471,10 +483,9 @@ def new_config_file_contents(
 
         conf[model] = stanza
 
-    # if no default model was chosen, then we select the first
-    # one in the list
+    # if no default model was chosen, then we select the first one in the list
     if not default_selected:
-        conf[list(successfully_downloaded.keys())[0]]["default"] = True
+        conf[list(conf.keys())[0]]["default"] = True
 
     return OmegaConf.to_yaml(conf)
 

ldm/generate.py

@@ -372,12 +372,6 @@ def _load_ckpt_model(self, model_name, mconfig):
         )
         from ldm.invoke.ckpt_to_diffuser import load_pipeline_from_original_stable_diffusion_ckpt
 
-#        try:
-#            if self.list_models()[self.current_model]['status'] == 'active':
-#                self.offload_model(self.current_model)
-#        except Exception:
-#            pass
-
         if self._has_cuda():
             torch.cuda.empty_cache()
         pipeline = load_pipeline_from_original_stable_diffusion_ckpt(
@@ -423,9 +417,9 @@ def _load_diffusers_model(self, mconfig):
             pipeline_args.update(cache_dir=global_cache_dir("hub"))
         if using_fp16:
             pipeline_args.update(torch_dtype=torch.float16)
-            fp_args_list = [{"revision": "fp16"}, {}]
-        else:
-            fp_args_list = [{}]
+        revision = mconfig.get('revision') or ('fp16' if using_fp16 else None)
+        fp_args_list = [{"revision": revision}] if revision else []
+        fp_args_list.append({})
 
         verbosity = dlogging.get_verbosity()
         dlogging.set_verbosity_error()