[Fix] Lora double hook (#3471)

Currently hooks registers multiple time for some modules.
As result - lora applies multiple time to this modules on generation and
images looks weird.
If have any other minds how to fix it better - feel free to push.

docs/features/LORAS.md

@@ -41,6 +41,16 @@ Windows systems). If the `loras` folder does not already exist, just
 create it. The vast majority of LoRA models use the Kohya file format,
 which is a type of `.safetensors` file.
 
+!!! warning "LoRA Naming Restrictions"
+
+    InvokeAI will only recognize LoRA files that contain the
+    characters a-z, A-Z, 0-9 and the underscore character
+    _. Other characters, including the hyphen, will cause the
+    LoRA file not to load. These naming restrictions may be
+    relaxed in the future, but for now you will need to rename
+    files that contain hyphens, commas, brackets, and other
+    non-word characters.
+
 You may change where InvokeAI looks for the `loras` folder by passing the
 `--lora_directory` option to the `invoke.sh`/`invoke.bat` launcher, or
 by placing the option in `invokeai.init`. For example:

docs/features/index.md

@@ -33,6 +33,11 @@ title: Overview
         Restore mangled faces and make images larger with upscaling. Also see
         the [Embiggen Upscaling Guide](EMBIGGEN.md).
 
+    -   The [Using LoRA Models](LORAS.md)
+
+        Add custom subjects and styles using HuggingFace's repository of
+        embeddings.
+	
     -   The [Concepts Library](CONCEPTS.md)
 
         Add custom subjects and styles using HuggingFace's repository of

docs/installation/deprecated_documentation/INSTALL_LINUX.md

@@ -79,7 +79,7 @@ title: Manual Installation, Linux
         and obtaining an access token for downloading. It will then download and
         install the weights files for you.
 
-        Please look [here](../INSTALL_MANUAL.md) for a manual process for doing
+        Please look [here](../020_INSTALL_MANUAL.md) for a manual process for doing
         the same thing.
 
 7.  Start generating images!

docs/installation/deprecated_documentation/INSTALL_WINDOWS.md

@@ -75,7 +75,7 @@ Note that you will need NVIDIA drivers, Python 3.10, and Git installed beforehan
           obtaining an access token for downloading. It will then download and install the
           weights files for you.
 
-          Please look [here](../INSTALL_MANUAL.md) for a manual process for doing the
+          Please look [here](../020_INSTALL_MANUAL.md) for a manual process for doing the
           same thing.
 
 8. Start generating images!

docs/requirements-mkdocs.txt

@@ -0,0 +1,5 @@
+mkdocs
+mkdocs-material>=8, <9
+mkdocs-git-revision-date-localized-plugin
+mkdocs-redirects==1.2.0
+

installer/lib/installer.py

@@ -243,16 +243,15 @@ class InvokeAiInstance:
 
         # Note that we're installing pinned versions of torch and
         # torchvision here, which *should* correspond to what is
-        # in pyproject.toml. This is to prevent torch 2.0 from
-        # being installed and immediately uninstalled and replaced with 1.13
+        # in pyproject.toml.
         pip = local[self.pip]
 
         (
             pip[
                 "install",
                 "--require-virtualenv",
-                "torch~=1.13.1",
-                "torchvision~=0.14.1",
+                "torch~=2.0.0",
+                "torchvision>=0.14.1",
                 "--force-reinstall",
                 "--find-links" if find_links is not None else None,
                 find_links,

invokeai/backend/invoke_ai_web_server.py

@@ -25,7 +25,7 @@ from invokeai.backend.modules.parameters import parameters_to_command
 import invokeai.frontend.dist as frontend
 from ldm.generate import Generate
 from ldm.invoke.args import Args, APP_ID, APP_VERSION, calculate_init_img_hash
-from ldm.invoke.concepts_lib import HuggingFaceConceptsLibrary
+from ldm.invoke.concepts_lib import get_hf_concepts_lib
 from ldm.invoke.conditioning import (
     get_tokens_for_prompt_object,
     get_prompt_structure,
@@ -538,7 +538,7 @@ class InvokeAIWebServer:
             try:
                 local_triggers = self.generate.model.textual_inversion_manager.get_all_trigger_strings()
                 locals = [{'name': x} for x in sorted(local_triggers, key=str.casefold)]
-                concepts = HuggingFaceConceptsLibrary().list_concepts(minimum_likes=5)
+                concepts = get_hf_concepts_lib().list_concepts(minimum_likes=5)
                 concepts = [{'name': f'<{x}>'} for x in sorted(concepts, key=str.casefold) if f'<{x}>' not in local_triggers]
                 socketio.emit("foundTextualInversionTriggers", {'local_triggers': locals, 'huggingface_concepts': concepts})
             except Exception as e:

ldm/generate.py

@@ -13,11 +13,16 @@ import time
 import traceback
 from typing import List
 
+import warnings
+with warnings.catch_warnings():
+    warnings.filterwarnings("ignore", category=UserWarning)
+    import torch
+    
 import cv2
 import diffusers
 import numpy as np
 import skimage
-import torch
+
 import transformers
 from diffusers.pipeline_utils import DiffusionPipeline
 from diffusers.utils.import_utils import is_xformers_available
@@ -979,13 +984,15 @@ class Generate:
         seed_everything(random.randrange(0, np.iinfo(np.uint32).max))
         if self.embedding_path and not model_data.get("ti_embeddings_loaded"):
             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
-                    )
-            model_data["ti_embeddings_loaded"] = True
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=UserWarning)
+                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
+                        )
+                model_data["ti_embeddings_loaded"] = True
         print(
             f'>> Textual inversion triggers: {", ".join(sorted(self.model.textual_inversion_manager.get_all_trigger_strings()))}'
         )

ldm/invoke/CLI.py

@@ -9,7 +9,6 @@ from pathlib import Path
 from typing import Union
 
 import click
-
 from compel import PromptParser
 
 if sys.platform == "darwin":

ldm/invoke/_version.py

@@ -1 +1,3 @@
-__version__='2.3.5-rc3'
+__version__='2.3.5.post2'
+
+

ldm/invoke/ckpt_to_diffuser.py

@@ -620,7 +620,10 @@ def convert_ldm_vae_checkpoint(checkpoint, config):
     for key in keys:
         if key.startswith(vae_key):
             vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key)
+    new_checkpoint = convert_ldm_vae_state_dict(vae_state_dict,config)
+    return new_checkpoint
 
+def convert_ldm_vae_state_dict(vae_state_dict, config):
     new_checkpoint = {}
 
     new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"]

ldm/invoke/concepts_lib.py

@@ -12,6 +12,14 @@ from urllib import request, error as ul_error
 from huggingface_hub import HfFolder, hf_hub_url, ModelSearchArguments, ModelFilter, HfApi
 from ldm.invoke.globals import Globals
 
+singleton = None
+
+def get_hf_concepts_lib():
+    global singleton
+    if singleton is None:
+        singleton = HuggingFaceConceptsLibrary()
+    return singleton
+
 class HuggingFaceConceptsLibrary(object):
     def __init__(self, root=None):
         '''

ldm/invoke/config/invokeai_update.py

@@ -6,6 +6,7 @@ import os
 import platform
 import psutil
 import requests
+import pkg_resources
 from rich import box, print
 from rich.console import Console, group
 from rich.panel import Panel
@@ -39,7 +40,7 @@ def invokeai_is_running()->bool:
             if matches:
                 print(f':exclamation: [bold red]An InvokeAI instance appears to be running as process {p.pid}[/red bold]')
                 return True
-        except psutil.AccessDenied:
+        except (psutil.AccessDenied,psutil.NoSuchProcess):
             continue
     return False
         
@@ -72,10 +73,20 @@ def welcome(versions: dict):
     )
     console.line()
 
+def get_extras():
+    extras = ''
+    try:
+        dist = pkg_resources.get_distribution('xformers')
+        extras = '[xformers]'
+    except pkg_resources.DistributionNotFound:
+        pass
+    return extras
+
 def main():
     versions = get_versions()
     if invokeai_is_running():
         print(f':exclamation: [bold red]Please terminate all running instances of InvokeAI before updating.[/red bold]')
+        input('Press any key to continue...')
         return
 
     welcome(versions)
@@ -94,13 +105,15 @@ def main():
     elif choice=='4':
         branch = Prompt.ask('Enter an InvokeAI branch name')
 
+    extras = get_extras()
+
     print(f':crossed_fingers: Upgrading to [yellow]{tag if tag else release}[/yellow]')
     if release:
-        cmd = f'pip install {INVOKE_AI_SRC}/{release}.zip --use-pep517 --upgrade'
+        cmd = f"pip install 'invokeai{extras} @ {INVOKE_AI_SRC}/{release}.zip' --use-pep517 --upgrade"
     elif tag:
-        cmd = f'pip install {INVOKE_AI_TAG}/{tag}.zip --use-pep517 --upgrade'
+        cmd = f"pip install 'invokeai{extras} @ {INVOKE_AI_TAG}/{tag}.zip' --use-pep517 --upgrade"
     else:
-        cmd = f'pip install {INVOKE_AI_BRANCH}/{branch}.zip --use-pep517 --upgrade'
+        cmd = f"pip install 'invokeai{extras} @ {INVOKE_AI_BRANCH}/{branch}.zip' --use-pep517 --upgrade"
     print('')
     print('')
     if os.system(cmd)==0:

ldm/invoke/model_manager.py

@@ -9,7 +9,6 @@ from __future__ import annotations
 import contextlib
 import gc
 import hashlib
-import io
 import os
 import re
 import sys
@@ -31,11 +30,10 @@ from huggingface_hub import scan_cache_dir
 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
-from ldm.util import ask_user, download_with_resume, instantiate_from_config, url_attachment_name
+from ldm.util import ask_user, download_with_resume, url_attachment_name
 
 
 class SDLegacyType(Enum):
@@ -370,8 +368,9 @@ class ModelManager(object):
         print(
             f">> Converting legacy checkpoint {model_name} into a diffusers model..."
         )
-        from ldm.invoke.ckpt_to_diffuser import load_pipeline_from_original_stable_diffusion_ckpt
-
+        from .ckpt_to_diffuser import (
+            load_pipeline_from_original_stable_diffusion_ckpt,
+        )
         if self._has_cuda():
             torch.cuda.empty_cache()
         pipeline = load_pipeline_from_original_stable_diffusion_ckpt(
@@ -433,7 +432,7 @@ class ModelManager(object):
                     **fp_args,
                 )
             except OSError as e:
-                if str(e).startswith("fp16 is not a valid"):
+                if 'Revision Not Found' in str(e):
                     pass
                 else:
                     print(
@@ -1230,6 +1229,17 @@ class ModelManager(object):
         return vae_path
                                
     def _load_vae(self, vae_config) -> AutoencoderKL:
+        using_fp16 = self.precision == "float16"
+        dtype = torch.float16 if using_fp16 else torch.float32
+        
+        # Handle the common case of a user shoving a VAE .ckpt into
+        # the vae field for a diffusers. We convert it into diffusers
+        # format and use it.
+        if isinstance(vae_config,(str,Path)):
+            return self.convert_vae(vae_config).to(dtype=dtype)
+        elif isinstance(vae_config,DictConfig) and (vae_path := vae_config.get('path')):
+            return self.convert_vae(vae_path).to(dtype=dtype)
+            
         vae_args = {}
         try:
             name_or_path = self.model_name_or_path(vae_config)
@@ -1237,7 +1247,6 @@ class ModelManager(object):
             return None
         if name_or_path is None:
             return None
-        using_fp16 = self.precision == "float16"
 
         vae_args.update(
             cache_dir=global_cache_dir("hub"),
@@ -1277,6 +1286,32 @@ class ModelManager(object):
 
         return vae
 
+    @staticmethod
+    def convert_vae(vae_path: Union[Path,str])->AutoencoderKL:
+        print("   | A checkpoint VAE was detected. Converting to diffusers format.")
+        vae_path = Path(Globals.root,vae_path).resolve()
+        
+        from .ckpt_to_diffuser import (
+            create_vae_diffusers_config,
+            convert_ldm_vae_state_dict,
+        )
+
+        vae_path = Path(vae_path)
+        if vae_path.suffix in ['.pt','.ckpt']:
+            vae_state_dict = torch.load(vae_path, map_location="cpu")
+        else:
+            vae_state_dict = safetensors.torch.load_file(vae_path)
+        if 'state_dict' in vae_state_dict:
+            vae_state_dict = vae_state_dict['state_dict']
+        # TODO: see if this works with 1.x inpaint models and 2.x models
+        config_file_path = Path(Globals.root,"configs/stable-diffusion/v1-inference.yaml")
+        original_conf = OmegaConf.load(config_file_path)
+        vae_config = create_vae_diffusers_config(original_conf, image_size=512) # TODO: fix
+        diffusers_vae = convert_ldm_vae_state_dict(vae_state_dict,vae_config)
+        vae = AutoencoderKL(**vae_config)
+        vae.load_state_dict(diffusers_vae)
+        return vae
+
     @staticmethod
     def _delete_model_from_cache(repo_id):
         cache_info = scan_cache_dir(global_cache_dir("diffusers"))

ldm/invoke/readline.py

@@ -13,7 +13,7 @@ import re
 import atexit
 from typing import List
 from ldm.invoke.args import Args
-from ldm.invoke.concepts_lib import HuggingFaceConceptsLibrary
+from ldm.invoke.concepts_lib import get_hf_concepts_lib
 from ldm.invoke.globals import Globals
 from ldm.modules.lora_manager import LoraManager
 
@@ -287,7 +287,7 @@ class Completer(object):
     def _concept_completions(self, text, state):
         if self.concepts is None:
             # cache Concepts() instance so we can check for updates in concepts_list during runtime.
-            self.concepts = HuggingFaceConceptsLibrary()
+            self.concepts = get_hf_concepts_lib()
             self.embedding_terms.update(set(self.concepts.list_concepts()))
         else:
             self.embedding_terms.update(set(self.concepts.list_concepts()))

ldm/modules/embedding_manager.py

@@ -6,7 +6,7 @@ from torch import nn
 
 import sys
 
-from ldm.invoke.concepts_lib import HuggingFaceConceptsLibrary
+from ldm.invoke.concepts_lib import get_hf_concepts_lib
 from ldm.data.personalized import per_img_token_list
 from transformers import CLIPTokenizer
 from functools import partial
@@ -39,7 +39,7 @@ class EmbeddingManager(nn.Module):
         super().__init__()
 
         self.embedder = embedder
-        self.concepts_library=HuggingFaceConceptsLibrary()
+        self.concepts_library=get_hf_concepts_lib()
 
         self.string_to_token_dict = {}
         self.string_to_param_dict = nn.ParameterDict()

ldm/modules/kohya_lora_manager.py

@@ -1,6 +1,6 @@
 import json
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Dict, Tuple
 
 import torch
 from diffusers.models import UNet2DConditionModel
@@ -166,12 +166,12 @@ class LoKRLayer:
 class LoRAModuleWrapper:
     unet: UNet2DConditionModel
     text_encoder: CLIPTextModel
-    hooks: list[RemovableHandle]
+    hooks: Dict[str, Tuple[torch.nn.Module, RemovableHandle]]
 
     def __init__(self, unet, text_encoder):
         self.unet = unet
         self.text_encoder = text_encoder
-        self.hooks = []
+        self.hooks = dict()
         self.text_modules = None
         self.unet_modules = None
 
@@ -228,7 +228,7 @@ class LoRAModuleWrapper:
         wrapper = self
 
         def lora_forward(module, input_h, output):
-            if len(wrapper.loaded_loras) == 0:
+            if len(wrapper.applied_loras) == 0:
                 return output
 
             for lora in wrapper.applied_loras.values():
@@ -241,11 +241,18 @@ class LoRAModuleWrapper:
         return lora_forward
 
     def apply_module_forward(self, module, name):
-        handle = module.register_forward_hook(self.lora_forward_hook(name))
-        self.hooks.append(handle)
+        if name in self.hooks:
+            registered_module, _ = self.hooks[name]
+            if registered_module != module:
+                raise Exception(f"Trying to register multiple modules to lora key: {name}")
+            # else it's just double hook creation - nothing to do
+            
+        else:
+            handle = module.register_forward_hook(self.lora_forward_hook(name))
+            self.hooks[name] = (module, handle)
 
     def clear_hooks(self):
-        for hook in self.hooks:
+        for _, hook in self.hooks.values():
             hook.remove()
 
         self.hooks.clear()

ldm/modules/textual_inversion_manager.py

@@ -3,14 +3,16 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional, Union
 
-import safetensors.torch
-import torch
+import warnings
+with warnings.catch_warnings():
+    warnings.filterwarnings("ignore", category=UserWarning)
+    import safetensors.torch
+    import torch
 from picklescan.scanner import scan_file_path
 from transformers import CLIPTextModel, CLIPTokenizer
 
 from compel.embeddings_provider import BaseTextualInversionManager
-from ldm.invoke.concepts_lib import HuggingFaceConceptsLibrary
-
+from ldm.invoke.concepts_lib import get_hf_concepts_lib
 
 @dataclass
 class TextualInversion:
@@ -34,7 +36,7 @@ class TextualInversionManager(BaseTextualInversionManager):
         self.tokenizer = tokenizer
         self.text_encoder = text_encoder
         self.full_precision = full_precision
-        self.hf_concepts_library = HuggingFaceConceptsLibrary()
+        self.hf_concepts_library = get_hf_concepts_lib()
         self.trigger_to_sourcefile = dict()
         default_textual_inversions: list[TextualInversion] = []
         self.textual_inversions = default_textual_inversions

pyproject.toml

@@ -32,9 +32,9 @@ dependencies = [
   "albumentations",
   "click",
   "clip_anytorch",
-  "compel~=1.1.0",
+  "compel~=1.1.5",
   "datasets",
-  "diffusers[torch]~=0.15.0",
+  "diffusers[torch]~=0.16.1",
   "dnspython==2.2.1",
   "einops",
   "eventlet",
@@ -76,7 +76,7 @@ dependencies = [
   "taming-transformers-rom1504",
   "test-tube>=0.7.5",
   "torch-fidelity",
-  "torch~=1.13.1",
+  "torch~=2.0.0",
   "torchmetrics",
   "torchvision>=0.14.1",
   "transformers~=4.26",
@@ -108,7 +108,7 @@ requires-python = ">=3.9, <3.11"
 "test" = ["pytest-cov", "pytest>6.0.0"]
 "xformers" = [
   "triton; sys_platform=='linux'",
-  "xformers~=0.0.16; sys_platform!='darwin'",
+  "xformers~=0.0.19; sys_platform!='darwin'",
 ]
 
 [project.scripts]