From 141be95c2cb4bd6611f36766e087777a4425a3eb Mon Sep 17 00:00:00 2001 From: Jordan Date: Sat, 18 Feb 2023 05:29:04 -0700 Subject: [PATCH 01/30] initial setup of lora support --- binary_installer/requirements.in | 2 +- ldm/generate.py | 7 ++++ ldm/invoke/generator/diffusers_pipeline.py | 3 ++ ldm/modules/lora_manager.py | 43 ++++++++++++++++++++++ pyproject.toml | 2 +- 5 files changed, 55 insertions(+), 2 deletions(-) create mode 100644 ldm/modules/lora_manager.py diff --git a/binary_installer/requirements.in b/binary_installer/requirements.in index 66e0618f789..5922b6f04e9 100644 --- a/binary_installer/requirements.in +++ b/binary_installer/requirements.in @@ -4,7 +4,7 @@ --trusted-host https://download.pytorch.org accelerate~=0.15 albumentations -diffusers[torch]~=0.11 +diffusers[torch]~=0.13 einops eventlet flask_cors diff --git a/ldm/generate.py b/ldm/generate.py index 1b07122628f..8e21f39f33c 100644 --- a/ldm/generate.py +++ b/ldm/generate.py @@ -457,6 +457,9 @@ def process_image(image,seed): self.sampler_name = sampler_name self._set_sampler() + if self.model.lora_manager: + prompt = self.model.lora_manager.configure_prompt(prompt) + # apply the concepts library to the prompt prompt = self.huggingface_concepts_library.replace_concepts_with_triggers( prompt, @@ -515,6 +518,9 @@ def process_image(image,seed): 'extractor':self.safety_feature_extractor } if self.safety_checker else None + if self.model.lora_manager: + self.model.lora_manager.load_lora() + results = generator.generate( prompt, iterations=iterations, @@ -927,6 +933,7 @@ def set_model(self,model_name): self.model_name = model_name self._set_sampler() # requires self.model_name to be set first + return self.model def load_huggingface_concepts(self, concepts:list[str]): diff --git a/ldm/invoke/generator/diffusers_pipeline.py b/ldm/invoke/generator/diffusers_pipeline.py index 5990eb42a17..82e7f9afc9d 100644 --- a/ldm/invoke/generator/diffusers_pipeline.py +++ b/ldm/invoke/generator/diffusers_pipeline.py @@ -28,6 +28,7 @@ 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 ldm.modules.lora_manager import LoraManager from ..offloading import LazilyLoadedModelGroup, FullyLoadedModelGroup, ModelGroup from ...models.diffusion.cross_attention_map_saving import AttentionMapSaver from ...modules.prompt_to_embeddings_converter import WeightedPromptFragmentsToEmbeddingsConverter @@ -292,6 +293,8 @@ def __init__( self.textual_inversion_manager = TextualInversionManager(tokenizer=self.tokenizer, text_encoder=self.text_encoder, full_precision=use_full_precision) + self.lora_manager = LoraManager(self.unet) + # InvokeAI's interface for text embeddings and whatnot self.prompt_fragments_to_embeddings_converter = WeightedPromptFragmentsToEmbeddingsConverter( tokenizer=self.tokenizer, diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py new file mode 100644 index 00000000000..05e3e3d8f04 --- /dev/null +++ b/ldm/modules/lora_manager.py @@ -0,0 +1,43 @@ +import re +from pathlib import Path + +from ldm.invoke.globals import global_models_dir +from diffusers.models import UNet2DConditionModel + +class LoraManager: + + def __init__(self, model: UNet2DConditionModel): + self.weights = {} + self.model = model + self.lora_path = Path(global_models_dir(), 'lora') + self.lora_match = re.compile(r"]+)>") + self.prompt = None + + def apply_lora_model(self, args): + args = args.split(':') + name = args[0] + path = Path(self.lora_path, name) + + if path.is_dir(): + print(f"loading lora: {path}") + self.model.load_attn_procs(path.absolute().as_posix()) + + if len(args) == 2: + self.weights[name] = args[1] + + def load_lora_from_prompt(self, prompt: str): + + for m in re.findall(self.lora_match, prompt): + self.apply_lora_model(m) + + def load_lora(self): + self.load_lora_from_prompt(self.prompt) + + def configure_prompt(self, prompt: str) -> str: + self.prompt = prompt + + def found(m): + return "" + + return re.sub(self.lora_match, found, prompt) + diff --git a/pyproject.toml b/pyproject.toml index f3dfa69b911..f6e428924d6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,7 @@ dependencies = [ "click", "clip_anytorch", # replacing "clip @ https://github.com/openai/CLIP/archive/eaa22acb90a5876642d0507623e859909230a52d.zip", "datasets", - "diffusers[torch]~=0.11", + "diffusers[torch]~=0.13", "dnspython==2.2.1", "einops", "eventlet", From afc8639c2514800e6de5cf906fa691d9f2c28ac9 Mon Sep 17 00:00:00 2001 From: Jordan Date: Sat, 18 Feb 2023 21:07:34 -0700 Subject: [PATCH 02/30] add pending support for safetensors with cloneofsimo/lora --- ldm/invoke/generator/diffusers_pipeline.py | 2 +- ldm/modules/lora_manager.py | 32 +++++++++++++++++----- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/ldm/invoke/generator/diffusers_pipeline.py b/ldm/invoke/generator/diffusers_pipeline.py index 82e7f9afc9d..25291649616 100644 --- a/ldm/invoke/generator/diffusers_pipeline.py +++ b/ldm/invoke/generator/diffusers_pipeline.py @@ -293,7 +293,7 @@ def __init__( self.textual_inversion_manager = TextualInversionManager(tokenizer=self.tokenizer, text_encoder=self.text_encoder, full_precision=use_full_precision) - self.lora_manager = LoraManager(self.unet) + self.lora_manager = LoraManager(self) # InvokeAI's interface for text embeddings and whatnot self.prompt_fragments_to_embeddings_converter = WeightedPromptFragmentsToEmbeddingsConverter( diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index 05e3e3d8f04..844c2a554f5 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -1,14 +1,14 @@ import re from pathlib import Path - from ldm.invoke.globals import global_models_dir -from diffusers.models import UNet2DConditionModel +from lora_diffusion import tune_lora_scale, patch_pipe + class LoraManager: - def __init__(self, model: UNet2DConditionModel): + def __init__(self, pipe): self.weights = {} - self.model = model + self.pipe = pipe self.lora_path = Path(global_models_dir(), 'lora') self.lora_match = re.compile(r"]+)>") self.prompt = None @@ -17,13 +17,31 @@ def apply_lora_model(self, args): args = args.split(':') name = args[0] path = Path(self.lora_path, name) + file = Path(path, "pytorch_lora_weights.bin") - if path.is_dir(): + if path.is_dir() and file.is_file(): print(f"loading lora: {path}") - self.model.load_attn_procs(path.absolute().as_posix()) - + self.pipe.unet.load_attn_procs(path.absolute().as_posix()) if len(args) == 2: self.weights[name] = args[1] + else: + # converting and saving in diffusers format + path_file = Path(self.lora_path, f'{name}.ckpt') + if Path(self.lora_path, f'{name}.safetensors').exists(): + path_file = Path(self.lora_path, f'{name}.safetensors') + + if path_file.is_file(): + print(f"loading lora: {path}") + patch_pipe( + self.pipe, + path_file.absolute().as_posix(), + patch_text=True, + patch_ti=True, + patch_unet=True, + ) + if len(args) == 2: + tune_lora_scale(self.pipe.unet, args[1]) + tune_lora_scale(self.pipe.text_encoder, args[1]) def load_lora_from_prompt(self, prompt: str): From 5a7145c48529a11999c226322fd8ecd8780725a9 Mon Sep 17 00:00:00 2001 From: Jordan Date: Sat, 18 Feb 2023 23:18:41 -0700 Subject: [PATCH 03/30] Create convert_lora.py --- scripts/convert_lora.py | 86 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 86 insertions(+) create mode 100644 scripts/convert_lora.py diff --git a/scripts/convert_lora.py b/scripts/convert_lora.py new file mode 100644 index 00000000000..b30890c7569 --- /dev/null +++ b/scripts/convert_lora.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python + +import re +from pathlib import Path +import torch +from safetensors.torch import load_file +import argparse +from diffusers import UNet2DConditionModel +from diffusers.pipelines.stable_diffusion.convert_from_ckpt import create_unet_diffusers_config +from omegaconf import OmegaConf +import requests + + +def parse_args(input_args=None): + parser = argparse.ArgumentParser(description="Convert kohya lora to diffusers") + parser.add_argument( + "--lora_file", + type=str, + default=None, + required=True, + help="Lora file to convert", + ) + parser.add_argument( + "--output_dir", + type=str, + default="models/lora", + help="The output directory where converted lora will be saved", + ) + + if input_args is not None: + args = parser.parse_args(input_args) + else: + args = parser.parse_args() + + return args + + +def replace_key_blocks(match_obj): + k = match_obj.groups() + + return f"{k[0]}.{k[1]}" + + +def replace_key_out(match_obj): + return f"to_out" + + +def replace_key_main(match_obj): + k = match_obj.groups() + block = re.sub(r"(.+)_(\d+)", replace_key_blocks, k[0]) + out = re.sub(r"to_out_(\d+)", replace_key_out, k[4]) + + return f"{block}.attentions.{k[1]}.transformer_blocks.{k[2]}.attn{k[3]}.processor.{out}_lora.{k[5]}" + + +def main(args): + response = requests.get( + "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml" + ) + original_config = OmegaConf.create(response.text) + + new_dict = dict() + lora_file = Path(args.lora_file) + + if lora_file.suffix == '.safetensors': + checkpoint = load_file(args.lora_file) + else: + checkpoint = torch.load(args.lora_file) + + for idx, key in enumerate(checkpoint): + check = re.compile(r"lora_unet_(.+)_attentions_(\d+)_transformer_blocks_(\d+)_attn(\d+)_(.+).lora_(.+)") + if check.match(key): + new_key = check.sub(replace_key_main, key) + new_dict[new_key] = checkpoint[key] + + unet_config = create_unet_diffusers_config(original_config, image_size=512) + unet = UNet2DConditionModel(**unet_config) + unet.load_attn_procs(new_dict) + + output_dir = Path(args.output_dir, lora_file.name.split('.')[0]) + unet.save_attn_procs(output_dir.absolute().as_posix()) + + +if __name__ == "__main__": + args = parse_args() + main(args) From 82e4d5aed2ece8ce696a013014c187c910ef7432 Mon Sep 17 00:00:00 2001 From: Jordan Date: Sun, 19 Feb 2023 17:33:24 -0700 Subject: [PATCH 04/30] change to new method to load safetensors --- ldm/modules/lora_manager.py | 92 ++++++++++++++++++++++++++++--------- scripts/convert_lora.py | 86 ---------------------------------- 2 files changed, 71 insertions(+), 107 deletions(-) delete mode 100644 scripts/convert_lora.py diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index 844c2a554f5..008864a4b00 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -1,7 +1,68 @@ import re from pathlib import Path from ldm.invoke.globals import global_models_dir -from lora_diffusion import tune_lora_scale, patch_pipe +import torch +from safetensors.torch import load_file + +# modified from script at https://github.com/huggingface/diffusers/pull/2403 +def merge_lora_into_pipe(pipeline, checkpoint_path, alpha): + # load LoRA weight from .safetensors + state_dict = load_file(checkpoint_path, device=torch.cuda.current_device()) + + visited = [] + + # directly update weight in diffusers model + for key in state_dict: + + # it is suggested to print out the key, it usually will be something like below + # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight" + + # as we have set the alpha beforehand, so just skip + if ".alpha" in key or key in visited: + continue + if "text" in key: + layer_infos = key.split(".")[0].split("lora_te" + "_")[-1].split("_") + curr_layer = pipeline.text_encoder + else: + layer_infos = key.split(".")[0].split("lora_unet" + "_")[-1].split("_") + curr_layer = pipeline.unet + + # find the target layer + temp_name = layer_infos.pop(0) + while len(layer_infos) > -1: + try: + curr_layer = curr_layer.__getattr__(temp_name) + if len(layer_infos) > 0: + temp_name = layer_infos.pop(0) + elif len(layer_infos) == 0: + break + except Exception: + if len(temp_name) > 0: + temp_name += "_" + layer_infos.pop(0) + else: + temp_name = layer_infos.pop(0) + + pair_keys = [] + if "lora_down" in key: + pair_keys.append(key.replace("lora_down", "lora_up")) + pair_keys.append(key) + else: + pair_keys.append(key) + pair_keys.append(key.replace("lora_up", "lora_down")) + + # update weight + if len(state_dict[pair_keys[0]].shape) == 4: + weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32) + weight_down = state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32) + curr_layer.weight.data += float(alpha) * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3) + else: + weight_up = state_dict[pair_keys[0]].to(torch.float32) + weight_down = state_dict[pair_keys[1]].to(torch.float32) + curr_layer.weight.data += float(alpha) * torch.mm(weight_up, weight_down) + + # update visited list + for item in pair_keys: + visited.append(item) class LoraManager: @@ -16,32 +77,21 @@ def __init__(self, pipe): def apply_lora_model(self, args): args = args.split(':') name = args[0] + path = Path(self.lora_path, name) file = Path(path, "pytorch_lora_weights.bin") if path.is_dir() and file.is_file(): - print(f"loading lora: {path}") + print(f"loading diffusers lora: {path}") self.pipe.unet.load_attn_procs(path.absolute().as_posix()) - if len(args) == 2: - self.weights[name] = args[1] else: - # converting and saving in diffusers format - path_file = Path(self.lora_path, f'{name}.ckpt') - if Path(self.lora_path, f'{name}.safetensors').exists(): - path_file = Path(self.lora_path, f'{name}.safetensors') - - if path_file.is_file(): - print(f"loading lora: {path}") - patch_pipe( - self.pipe, - path_file.absolute().as_posix(), - patch_text=True, - patch_ti=True, - patch_unet=True, - ) - if len(args) == 2: - tune_lora_scale(self.pipe.unet, args[1]) - tune_lora_scale(self.pipe.text_encoder, args[1]) + file = Path(self.lora_path, f"{name}.safetensors") + print(f"loading lora: {file}") + alpha = 1 + if len(args) == 2: + alpha = args[1] + + merge_lora_into_pipe(self.pipe, file.absolute().as_posix(), alpha) def load_lora_from_prompt(self, prompt: str): diff --git a/scripts/convert_lora.py b/scripts/convert_lora.py deleted file mode 100644 index b30890c7569..00000000000 --- a/scripts/convert_lora.py +++ /dev/null @@ -1,86 +0,0 @@ -#!/usr/bin/env python - -import re -from pathlib import Path -import torch -from safetensors.torch import load_file -import argparse -from diffusers import UNet2DConditionModel -from diffusers.pipelines.stable_diffusion.convert_from_ckpt import create_unet_diffusers_config -from omegaconf import OmegaConf -import requests - - -def parse_args(input_args=None): - parser = argparse.ArgumentParser(description="Convert kohya lora to diffusers") - parser.add_argument( - "--lora_file", - type=str, - default=None, - required=True, - help="Lora file to convert", - ) - parser.add_argument( - "--output_dir", - type=str, - default="models/lora", - help="The output directory where converted lora will be saved", - ) - - if input_args is not None: - args = parser.parse_args(input_args) - else: - args = parser.parse_args() - - return args - - -def replace_key_blocks(match_obj): - k = match_obj.groups() - - return f"{k[0]}.{k[1]}" - - -def replace_key_out(match_obj): - return f"to_out" - - -def replace_key_main(match_obj): - k = match_obj.groups() - block = re.sub(r"(.+)_(\d+)", replace_key_blocks, k[0]) - out = re.sub(r"to_out_(\d+)", replace_key_out, k[4]) - - return f"{block}.attentions.{k[1]}.transformer_blocks.{k[2]}.attn{k[3]}.processor.{out}_lora.{k[5]}" - - -def main(args): - response = requests.get( - "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml" - ) - original_config = OmegaConf.create(response.text) - - new_dict = dict() - lora_file = Path(args.lora_file) - - if lora_file.suffix == '.safetensors': - checkpoint = load_file(args.lora_file) - else: - checkpoint = torch.load(args.lora_file) - - for idx, key in enumerate(checkpoint): - check = re.compile(r"lora_unet_(.+)_attentions_(\d+)_transformer_blocks_(\d+)_attn(\d+)_(.+).lora_(.+)") - if check.match(key): - new_key = check.sub(replace_key_main, key) - new_dict[new_key] = checkpoint[key] - - unet_config = create_unet_diffusers_config(original_config, image_size=512) - unet = UNet2DConditionModel(**unet_config) - unet.load_attn_procs(new_dict) - - output_dir = Path(args.output_dir, lora_file.name.split('.')[0]) - unet.save_attn_procs(output_dir.absolute().as_posix()) - - -if __name__ == "__main__": - args = parse_args() - main(args) From 096e1d3a5d36fc2cc5cf5730e02b5f635d8df04d Mon Sep 17 00:00:00 2001 From: Jordan Date: Mon, 20 Feb 2023 02:37:44 -0700 Subject: [PATCH 05/30] start of rewrite for add / remove --- ldm/generate.py | 9 ++ ldm/modules/lora_manager.py | 238 ++++++++++++++++++++++++++++-------- 2 files changed, 195 insertions(+), 52 deletions(-) diff --git a/ldm/generate.py b/ldm/generate.py index 8e21f39f33c..28337874586 100644 --- a/ldm/generate.py +++ b/ldm/generate.py @@ -575,7 +575,13 @@ def process_image(image,seed): save_original = save_original, image_callback = image_callback) + if self.model.lora_manager: + self.model.lora_manager.reset_lora() + except KeyboardInterrupt: + if self.model.lora_manager: + self.model.lora_manager.reset_lora() + # Clear the CUDA cache on an exception self.clear_cuda_cache() @@ -584,6 +590,9 @@ def process_image(image,seed): else: raise KeyboardInterrupt except RuntimeError: + if self.model.lora_manager: + self.model.lora_manager.reset_lora() + # Clear the CUDA cache on an exception self.clear_cuda_cache() diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index 008864a4b00..16b1d4363c5 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -3,44 +3,99 @@ from ldm.invoke.globals import global_models_dir import torch from safetensors.torch import load_file +from typing import List, Optional, Set, Type -# modified from script at https://github.com/huggingface/diffusers/pull/2403 -def merge_lora_into_pipe(pipeline, checkpoint_path, alpha): - # load LoRA weight from .safetensors - state_dict = load_file(checkpoint_path, device=torch.cuda.current_device()) - visited = [] +class LoraLinear(torch.nn.Module): + def __init__( + self, in_features, out_features, rank=4 + ): + super().__init__() - # directly update weight in diffusers model - for key in state_dict: + if rank > min(in_features, out_features): + raise ValueError( + f"LoRA rank {rank} must be less or equal than {min(in_features, out_features)}" + ) + self.rank = rank + self.linear = torch.nn.Linear(in_features, out_features, bias=False) + self.lora = torch.nn.Linear(in_features, out_features, bias=False) - # it is suggested to print out the key, it usually will be something like below - # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight" + def forward(self, hidden_states): + orig_dtype = hidden_states.dtype + dtype = self.lora.weight.dtype + return self.lora(hidden_states.to(dtype)).to(orig_dtype) - # as we have set the alpha beforehand, so just skip - if ".alpha" in key or key in visited: - continue - if "text" in key: - layer_infos = key.split(".")[0].split("lora_te" + "_")[-1].split("_") - curr_layer = pipeline.text_encoder + +class LoraManager: + + def __init__(self, pipe): + self.pipe = pipe + self.lora_path = Path(global_models_dir(), 'lora') + self.lora_match = re.compile(r"]+)>") + self.prompt = None + + def _process_lora(self, lora): + processed_lora = { + "unet": [], + "text_encoder": [] + } + visited = [] + for key in lora: + if ".alpha" in key or key in visited: + continue + if "text" in key: + lora_type, pair_keys = self._find_layer( + "text_encoder", + key.split(".")[0].split("lora_te" + "_")[-1].split("_"), + key + ) + else: + lora_type, pair_keys = self._find_layer( + "unet", + key.split(".")[0].split("lora_unet" + "_")[-1].split("_"), + key + ) + + if len(lora[pair_keys[0]].shape) == 4: + weight_up = lora[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32) + weight_down = lora[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32) + weight = torch.mm(weight_up, weight_down) + else: + weight_up = lora[pair_keys[0]].to(torch.float32) + weight_down = lora[pair_keys[1]].to(torch.float32) + weight = torch.mm(weight_up, weight_down) + + processed_lora[lora_type].append({ + "weight": weight, + "rank": lora[pair_keys[1]].shape[0] + }) + + for item in pair_keys: + visited.append(item) + + return processed_lora + + def _find_layer(self, lora_type, layer_key, key): + temp_name = layer_key.pop(0) + if lora_type == "unet": + curr_layer = self.pipe.unet + elif lora_type == "text_encoder": + curr_layer = self.pipe.text_encoder else: - layer_infos = key.split(".")[0].split("lora_unet" + "_")[-1].split("_") - curr_layer = pipeline.unet + raise ValueError("Invalid Lora Type") - # find the target layer - temp_name = layer_infos.pop(0) - while len(layer_infos) > -1: + while len(layer_key) > -1: try: curr_layer = curr_layer.__getattr__(temp_name) - if len(layer_infos) > 0: - temp_name = layer_infos.pop(0) - elif len(layer_infos) == 0: + if len(layer_key) > 0: + temp_name = layer_key.pop(0) + elif len(layer_key) == 0: break except Exception: if len(temp_name) > 0: - temp_name += "_" + layer_infos.pop(0) + temp_name += "_" + layer_key.pop(0) else: - temp_name = layer_infos.pop(0) + temp_name = layer_key.pop(0) pair_keys = [] if "lora_down" in key: @@ -50,29 +105,89 @@ def merge_lora_into_pipe(pipeline, checkpoint_path, alpha): pair_keys.append(key) pair_keys.append(key.replace("lora_up", "lora_down")) - # update weight - if len(state_dict[pair_keys[0]].shape) == 4: - weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32) - weight_down = state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32) - curr_layer.weight.data += float(alpha) * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3) + return lora_type, pair_keys + + @staticmethod + def _find_modules( + model, + ancestor_class: Optional[Set[str]] = None, + search_class: List[Type[torch.nn.Module]] = [torch.nn.Linear], + exclude_children_of: Optional[List[Type[torch.nn.Module]]] = [LoraLinear], + ): + """ + Find all modules of a certain class (or union of classes) that are direct or + indirect descendants of other modules of a certain class (or union of classes). + Returns all matching modules, along with the parent of those modules and the + names they are referenced by. + """ + + # Get the targets we should replace all linears under + if ancestor_class is not None: + ancestors = ( + module + for module in model.modules() + if module.__class__.__name__ in ancestor_class + ) else: - weight_up = state_dict[pair_keys[0]].to(torch.float32) - weight_down = state_dict[pair_keys[1]].to(torch.float32) - curr_layer.weight.data += float(alpha) * torch.mm(weight_up, weight_down) - - # update visited list - for item in pair_keys: - visited.append(item) - - -class LoraManager: - - def __init__(self, pipe): - self.weights = {} - self.pipe = pipe - self.lora_path = Path(global_models_dir(), 'lora') - self.lora_match = re.compile(r"]+)>") - self.prompt = None + # this, incase you want to naively iterate over all modules. + ancestors = [module for module in model.modules()] + + # For each target find every linear_class module that isn't a child of a LoraInjectedLinear + for ancestor in ancestors: + for fullname, module in ancestor.named_modules(): + if any([isinstance(module, _class) for _class in search_class]): + # Find the direct parent if this is a descendant, not a child, of target + *path, name = fullname.split(".") + parent = ancestor + while path: + parent = parent.get_submodule(path.pop(0)) + # Skip this linear if it's a child of a LoraInjectedLinear + if exclude_children_of and any( + [isinstance(parent, _class) for _class in exclude_children_of] + ): + continue + # Otherwise, yield it + yield parent, name, module + + @staticmethod + def patch_module(lora_type, processed_lora, module, name, child_module, scale: float = 1.0): + _source = ( + child_module.linear + if isinstance(child_module, LoraLinear) + else child_module + ) + + lora = processed_lora[lora_type].pop(0) + + weight = _source.weight + _tmp = LoraLinear( + in_features=_source.in_features, + out_features=_source.out_features, + rank=lora["rank"] + ) + _tmp.linear.weight = weight + + # switch the module + module._modules[name] = _tmp + module._modules[name].lora.weight.data = lora["weight"] + module._modules[name].to(weight.device) + + def patch_lora(self, lora_path, scale: float = 1.0): + lora = load_file(lora_path) + processed_lora = self._process_lora(lora) + for module, name, child_module in self._find_modules( + self.pipe.unet, + {"CrossAttention", "Attention", "GEGLU"}, + search_class=[torch.nn.Linear, LoraLinear] + ): + self.patch_module("unet", processed_lora, module, name, child_module, scale) + + for module, name, child_module in self._find_modules( + self.pipe.text_encoder, + {"CLIPAttention"}, + search_class=[torch.nn.Linear, LoraLinear] + ): + self.patch_module("text_encoder", processed_lora, module, name, child_module, scale) def apply_lora_model(self, args): args = args.split(':') @@ -87,14 +202,34 @@ def apply_lora_model(self, args): else: file = Path(self.lora_path, f"{name}.safetensors") print(f"loading lora: {file}") - alpha = 1 + scale = 1.0 if len(args) == 2: - alpha = args[1] + scale = float(args[1]) - merge_lora_into_pipe(self.pipe, file.absolute().as_posix(), alpha) + self.patch_lora(file.absolute().as_posix(), scale) - def load_lora_from_prompt(self, prompt: str): + @staticmethod + def remove_lora(child_module): + _source = child_module.linear + weight = _source.weight + + _tmp = torch.nn.Linear(_source.in_features, _source.out_features) + _tmp.weight = weight + + def reset_lora(self): + for module, name, child_module in self._find_modules( + self.pipe.unet, + search_class=[LoraLinear] + ): + self.remove_lora(child_module) + for module, name, child_module in self._find_modules( + self.pipe.text_encoder, + search_class=[LoraLinear] + ): + self.remove_lora(child_module) + + def load_lora_from_prompt(self, prompt: str): for m in re.findall(self.lora_match, prompt): self.apply_lora_model(m) @@ -108,4 +243,3 @@ def found(m): return "" return re.sub(self.lora_match, found, prompt) - From e744774171f7f58c8e626c441521d25053cd50de Mon Sep 17 00:00:00 2001 From: neecapp Date: Mon, 20 Feb 2023 13:49:16 -0500 Subject: [PATCH 06/30] Rewrite lora manager with hooks --- ldm/modules/lora_manager.py | 387 +++++++++++++++++------------------- 1 file changed, 186 insertions(+), 201 deletions(-) diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index 16b1d4363c5..8b3784a1e7e 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -1,241 +1,217 @@ import re from pathlib import Path from ldm.invoke.globals import global_models_dir -import torch +from ldm.invoke.devices import choose_torch_device from safetensors.torch import load_file -from typing import List, Optional, Set, Type +import torch +from torch.utils.hooks import RemovableHandle + +class LoRALayer: + lora_name: str + name: str + scale: float + up: torch.nn.Module + down: torch.nn.Module + + def __init__(self, lora_name: str, name: str, rank=4, alpha=1.0): + self.lora_name = lora_name + self.name = name + self.scale = alpha / rank + + +class LoRA: + name: str + layers: dict[str, LoRALayer] + multiplier: float + + def __init__(self, name: str, multiplier=1.0): + self.name = name + self.layers = {} + self.multiplier = multiplier + + +UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"] +TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"] +LORA_PREFIX_UNET = 'lora_unet' +LORA_PREFIX_TEXT_ENCODER = 'lora_te' + + +def load_lora( + name: str, + path_file: Path, + device: torch.device, + dtype: torch.dtype, + text_modules: dict[str, torch.nn.Module], + unet_modules: dict[str, torch.nn.Module], + multiplier=1.0 +): + print(f">> Loading lora {name} from {path_file}") + if path_file.suffix == '.safetensors': + checkpoint = load_file(path_file, device='cpu') + else: + checkpoint = torch.load(path_file, map_location='cpu') + + lora = LoRA(name, multiplier) + + alpha = None + rank = None + for key, value in checkpoint.items(): + stem, leaf = key.split(".", 1) + + if leaf.endswith("alpha"): + if alpha is None: + alpha = value.item() + continue + + if stem.startswith(LORA_PREFIX_TEXT_ENCODER): + # text encoder layer + wrapped = text_modules.get(stem, None) + if wrapped is None: + print(f">> Missing layer: {stem}") + continue + elif stem.startswith(LORA_PREFIX_UNET): + # unet layer + wrapped = unet_modules.get(stem, None) + if wrapped is None: + print(f">> Missing layer: {stem}") + continue + else: + continue + + if rank is None and leaf == 'lora_down.weight' and len(value.size()) == 2: + rank = value.shape[0] + if wrapped is None: + continue -class LoraLinear(torch.nn.Module): - def __init__( - self, in_features, out_features, rank=4 - ): - super().__init__() + layer = lora.layers.get(stem, None) + if layer is None: + layer = LoRALayer(name, stem, rank, alpha) + lora.layers[stem] = layer - if rank > min(in_features, out_features): - raise ValueError( - f"LoRA rank {rank} must be less or equal than {min(in_features, out_features)}" - ) - self.rank = rank - self.linear = torch.nn.Linear(in_features, out_features, bias=False) - self.lora = torch.nn.Linear(in_features, out_features, bias=False) + if type(wrapped) == torch.nn.Linear: + module = torch.nn.Linear( + value.shape[1], value.shape[0], bias=False) + elif type(wrapped) == torch.nn.Conv2d: + module = torch.nn.Conv2d( + value.shape[1], value.shape[0], (1, 1), bias=False) + else: + print( + f">> Encoundered unknown lora layer module in {name}: {type(value).__name__}") + + with torch.no_grad(): + module.weight.copy_(value) + + module.to(device=device, dtype=dtype) + + if leaf == "lora_up.weight": + layer.up = module + elif leaf == "lora_down.weight": + layer.down = module + else: + print(f">> Encountered unknown layer in lora {name}: {key}") + continue - def forward(self, hidden_states): - orig_dtype = hidden_states.dtype - dtype = self.lora.weight.dtype - return self.lora(hidden_states.to(dtype)).to(orig_dtype) + return lora class LoraManager: + loras: dict[str, LoRA] + applied_loras: dict[str, LoRA] + hooks: list[RemovableHandle] def __init__(self, pipe): - self.pipe = pipe self.lora_path = Path(global_models_dir(), 'lora') self.lora_match = re.compile(r"]+)>") - self.prompt = None - - def _process_lora(self, lora): - processed_lora = { - "unet": [], - "text_encoder": [] - } - visited = [] - for key in lora: - if ".alpha" in key or key in visited: - continue - if "text" in key: - lora_type, pair_keys = self._find_layer( - "text_encoder", - key.split(".")[0].split("lora_te" + "_")[-1].split("_"), - key - ) - else: - lora_type, pair_keys = self._find_layer( - "unet", - key.split(".")[0].split("lora_unet" + "_")[-1].split("_"), - key - ) - - if len(lora[pair_keys[0]].shape) == 4: - weight_up = lora[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32) - weight_down = lora[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32) - weight = torch.mm(weight_up, weight_down) - else: - weight_up = lora[pair_keys[0]].to(torch.float32) - weight_down = lora[pair_keys[1]].to(torch.float32) - weight = torch.mm(weight_up, weight_down) - - processed_lora[lora_type].append({ - "weight": weight, - "rank": lora[pair_keys[1]].shape[0] - }) - - for item in pair_keys: - visited.append(item) - - return processed_lora - - def _find_layer(self, lora_type, layer_key, key): - temp_name = layer_key.pop(0) - if lora_type == "unet": - curr_layer = self.pipe.unet - elif lora_type == "text_encoder": - curr_layer = self.pipe.text_encoder - else: - raise ValueError("Invalid Lora Type") - - while len(layer_key) > -1: - try: - curr_layer = curr_layer.__getattr__(temp_name) - if len(layer_key) > 0: - temp_name = layer_key.pop(0) - elif len(layer_key) == 0: - break - except Exception: - if len(temp_name) > 0: - temp_name += "_" + layer_key.pop(0) - else: - temp_name = layer_key.pop(0) - - pair_keys = [] - if "lora_down" in key: - pair_keys.append(key.replace("lora_down", "lora_up")) - pair_keys.append(key) - else: - pair_keys.append(key) - pair_keys.append(key.replace("lora_up", "lora_down")) - - return lora_type, pair_keys - - @staticmethod - def _find_modules( - model, - ancestor_class: Optional[Set[str]] = None, - search_class: List[Type[torch.nn.Module]] = [torch.nn.Linear], - exclude_children_of: Optional[List[Type[torch.nn.Module]]] = [LoraLinear], - ): - """ - Find all modules of a certain class (or union of classes) that are direct or - indirect descendants of other modules of a certain class (or union of classes). - Returns all matching modules, along with the parent of those modules and the - names they are referenced by. - """ - - # Get the targets we should replace all linears under - if ancestor_class is not None: - ancestors = ( - module - for module in model.modules() - if module.__class__.__name__ in ancestor_class - ) - else: - # this, incase you want to naively iterate over all modules. - ancestors = [module for module in model.modules()] - - # For each target find every linear_class module that isn't a child of a LoraInjectedLinear - for ancestor in ancestors: - for fullname, module in ancestor.named_modules(): - if any([isinstance(module, _class) for _class in search_class]): - # Find the direct parent if this is a descendant, not a child, of target - *path, name = fullname.split(".") - parent = ancestor - while path: - parent = parent.get_submodule(path.pop(0)) - # Skip this linear if it's a child of a LoraInjectedLinear - if exclude_children_of and any( - [isinstance(parent, _class) for _class in exclude_children_of] - ): - continue - # Otherwise, yield it - yield parent, name, module - - @staticmethod - def patch_module(lora_type, processed_lora, module, name, child_module, scale: float = 1.0): - _source = ( - child_module.linear - if isinstance(child_module, LoraLinear) - else child_module - ) - - lora = processed_lora[lora_type].pop(0) - - weight = _source.weight - _tmp = LoraLinear( - in_features=_source.in_features, - out_features=_source.out_features, - rank=lora["rank"] - ) - _tmp.linear.weight = weight - - # switch the module - module._modules[name] = _tmp - module._modules[name].lora.weight.data = lora["weight"] - module._modules[name].to(weight.device) - - def patch_lora(self, lora_path, scale: float = 1.0): - lora = load_file(lora_path) - processed_lora = self._process_lora(lora) - for module, name, child_module in self._find_modules( - self.pipe.unet, - {"CrossAttention", "Attention", "GEGLU"}, - search_class=[torch.nn.Linear, LoraLinear] - ): - self.patch_module("unet", processed_lora, module, name, child_module, scale) - - for module, name, child_module in self._find_modules( - self.pipe.text_encoder, - {"CLIPAttention"}, - search_class=[torch.nn.Linear, LoraLinear] - ): - self.patch_module("text_encoder", processed_lora, module, name, child_module, scale) + self.unet = pipe.unet + self.text_encoder = pipe.text_encoder + self.device = torch.device(choose_torch_device()) + self.dtype = pipe.unet.dtype + self.loras = {} + self.applied_loras = {} + self.hooks = [] + + def find_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> dict[str, torch.nn.Module]: + mapping = {} + for name, module in root_module.named_modules(): + if module.__class__.__name__ in target_replace_modules: + for child_name, child_module in module.named_modules(): + if child_module.__class__.__name__ == "Linear" or (child_module.__class__.__name__ == "Conv2d" and child_module.kernel_size == (1, 1)): + # It's easier to just convert into "lora" naming instead of trying to inverse map from "lora" -> diffuser + lora_name = prefix + '.' + name + '.' + child_name + lora_name = lora_name.replace('.', '_') + mapping[lora_name] = child_module + self.hooks.append(child_module.register_forward_hook(self._make_hook(lora_name))) + return mapping + + self.text_modules = find_modules( + LORA_PREFIX_TEXT_ENCODER, self.text_encoder, TEXT_ENCODER_TARGET_REPLACE_MODULE) + self.unet_modules = find_modules( + LORA_PREFIX_UNET, self.unet, UNET_TARGET_REPLACE_MODULE) + + def _make_hook(self, layer: str): + def hook(module, input, output): + for lora in self.applied_loras.values(): + lora_layer = lora.layers.get(layer, None) + if lora_layer is None: + continue + output = output + \ + lora_layer.up(lora_layer.down(*input)) * \ + lora.multiplier * lora_layer.scale + return output + return hook + + def _load_lora(self, name, path_file, multiplier=1.0): + lora = load_lora(name, path_file, self.device, self.dtype, + self.text_modules, self.unet_modules, multiplier) + self.loras[name] = lora + self.applied_loras[name] = lora + return lora def apply_lora_model(self, args): args = args.split(':') name = args[0] - path = Path(self.lora_path, name) file = Path(path, "pytorch_lora_weights.bin") if path.is_dir() and file.is_file(): - print(f"loading diffusers lora: {path}") + print(f"loading lora: {path}") self.pipe.unet.load_attn_procs(path.absolute().as_posix()) - else: - file = Path(self.lora_path, f"{name}.safetensors") - print(f"loading lora: {file}") - scale = 1.0 if len(args) == 2: - scale = float(args[1]) - - self.patch_lora(file.absolute().as_posix(), scale) + self.weights[name] = float(args[1]) + else: + # converting and saving in diffusers format + path_file = Path(self.lora_path, f'{name}.ckpt') + if Path(self.lora_path, f'{name}.safetensors').exists(): + path_file = Path(self.lora_path, f'{name}.safetensors') - @staticmethod - def remove_lora(child_module): - _source = child_module.linear - weight = _source.weight + if not path_file.exists(): + print(f">> Unable to find lora: {name}") + return - _tmp = torch.nn.Linear(_source.in_features, _source.out_features) - _tmp.weight = weight + mult = 1.0 + if len(args) == 2: + mult = float(args[1]) - def reset_lora(self): - for module, name, child_module in self._find_modules( - self.pipe.unet, - search_class=[LoraLinear] - ): - self.remove_lora(child_module) + lora = self.loras.get(name, None) + if lora is None: + lora = self._load_lora(name, path_file, mult) - for module, name, child_module in self._find_modules( - self.pipe.text_encoder, - search_class=[LoraLinear] - ): - self.remove_lora(child_module) + lora.multiplier = mult + self.applied_loras[name] = lora def load_lora_from_prompt(self, prompt: str): + self.applied_loras = {} for m in re.findall(self.lora_match, prompt): self.apply_lora_model(m) def load_lora(self): self.load_lora_from_prompt(self.prompt) + def unload_lora(self, lora_name: str): + if lora_name in self.loras: + del self.loras[lora_name] + def configure_prompt(self, prompt: str) -> str: self.prompt = prompt @@ -243,3 +219,12 @@ def found(m): return "" return re.sub(self.lora_match, found, prompt) + + def __del__(self): + del self.loras + del self.applied_loras + del self.text_modules + del self.unet_modules + for cb in self.hooks: + cb.remove() + del self.hooks From 8f6e43d4a460c90db9c42dd8c01c53e51afe6a9e Mon Sep 17 00:00:00 2001 From: Jordan Date: Mon, 20 Feb 2023 14:06:58 -0700 Subject: [PATCH 07/30] code cleanup --- ldm/generate.py | 9 --------- ldm/modules/lora_manager.py | 9 +++++---- 2 files changed, 5 insertions(+), 13 deletions(-) diff --git a/ldm/generate.py b/ldm/generate.py index 28337874586..8e21f39f33c 100644 --- a/ldm/generate.py +++ b/ldm/generate.py @@ -575,13 +575,7 @@ def process_image(image,seed): save_original = save_original, image_callback = image_callback) - if self.model.lora_manager: - self.model.lora_manager.reset_lora() - except KeyboardInterrupt: - if self.model.lora_manager: - self.model.lora_manager.reset_lora() - # Clear the CUDA cache on an exception self.clear_cuda_cache() @@ -590,9 +584,6 @@ def process_image(image,seed): else: raise KeyboardInterrupt except RuntimeError: - if self.model.lora_manager: - self.model.lora_manager.reset_lora() - # Clear the CUDA cache on an exception self.clear_cuda_cache() diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index 8b3784a1e7e..d21e179a8b8 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -6,6 +6,7 @@ import torch from torch.utils.hooks import RemovableHandle + class LoRALayer: lora_name: str name: str @@ -47,7 +48,7 @@ def load_lora( ): print(f">> Loading lora {name} from {path_file}") if path_file.suffix == '.safetensors': - checkpoint = load_file(path_file, device='cpu') + checkpoint = load_file(path_file.absolute().as_posix(), device='cpu') else: checkpoint = torch.load(path_file, map_location='cpu') @@ -98,6 +99,7 @@ def load_lora( else: print( f">> Encoundered unknown lora layer module in {name}: {type(value).__name__}") + continue with torch.no_grad(): module.weight.copy_(value) @@ -130,6 +132,7 @@ def __init__(self, pipe): self.loras = {} self.applied_loras = {} self.hooks = [] + self.prompt = "" def find_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> dict[str, torch.nn.Module]: mapping = {} @@ -176,9 +179,7 @@ def apply_lora_model(self, args): if path.is_dir() and file.is_file(): print(f"loading lora: {path}") - self.pipe.unet.load_attn_procs(path.absolute().as_posix()) - if len(args) == 2: - self.weights[name] = float(args[1]) + self.unet.load_attn_procs(path.absolute().as_posix()) else: # converting and saving in diffusers format path_file = Path(self.lora_path, f'{name}.ckpt') From 3c6c18b34c745647d1675c5335af8eb3108c9a1b Mon Sep 17 00:00:00 2001 From: Jordan Date: Mon, 20 Feb 2023 15:19:29 -0700 Subject: [PATCH 08/30] cleanup suggestions from neecap --- ldm/modules/lora_manager.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index d21e179a8b8..d33d15d961e 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -168,7 +168,6 @@ def _load_lora(self, name, path_file, multiplier=1.0): lora = load_lora(name, path_file, self.device, self.dtype, self.text_modules, self.unet_modules, multiplier) self.loras[name] = lora - self.applied_loras[name] = lora return lora def apply_lora_model(self, args): @@ -202,7 +201,6 @@ def apply_lora_model(self, args): self.applied_loras[name] = lora def load_lora_from_prompt(self, prompt: str): - self.applied_loras = {} for m in re.findall(self.lora_match, prompt): self.apply_lora_model(m) @@ -214,6 +212,7 @@ def unload_lora(self, lora_name: str): del self.loras[lora_name] def configure_prompt(self, prompt: str) -> str: + self.applied_loras = {} self.prompt = prompt def found(m): From ac972ebbe353edc6e4977b9a4eb8c33e900c30f2 Mon Sep 17 00:00:00 2001 From: Jordan Date: Mon, 20 Feb 2023 16:06:30 -0700 Subject: [PATCH 09/30] update prompt setup so lora's can be loaded in other ways --- ldm/modules/lora_manager.py | 39 ++++++++++++++++++++++++------------- 1 file changed, 25 insertions(+), 14 deletions(-) diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index d33d15d961e..d24316ddc6c 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -132,7 +132,7 @@ def __init__(self, pipe): self.loras = {} self.applied_loras = {} self.hooks = [] - self.prompt = "" + self.loras_to_load = [] def find_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> dict[str, torch.nn.Module]: mapping = {} @@ -170,9 +170,7 @@ def _load_lora(self, name, path_file, multiplier=1.0): self.loras[name] = lora return lora - def apply_lora_model(self, args): - args = args.split(':') - name = args[0] + def apply_lora_model(self, name, mult: float = 1.0): path = Path(self.lora_path, name) file = Path(path, "pytorch_lora_weights.bin") @@ -189,10 +187,6 @@ def apply_lora_model(self, args): print(f">> Unable to find lora: {name}") return - mult = 1.0 - if len(args) == 2: - mult = float(args[1]) - lora = self.loras.get(name, None) if lora is None: lora = self._load_lora(name, path_file, mult) @@ -200,20 +194,36 @@ def apply_lora_model(self, args): lora.multiplier = mult self.applied_loras[name] = lora - def load_lora_from_prompt(self, prompt: str): - for m in re.findall(self.lora_match, prompt): - self.apply_lora_model(m) - def load_lora(self): - self.load_lora_from_prompt(self.prompt) + for lora_to_load in self.loras_to_load: + self.apply_lora_model(lora_to_load["name"], lora_to_load["mult"]) def unload_lora(self, lora_name: str): if lora_name in self.loras: del self.loras[lora_name] + def set_lora(self, name, mult: float = 1.0): + if name in self.loras_to_load: + index = self.loras_to_load.index(name) + self.loras_to_load[index]["mult"] = mult + else: + self.loras_to_load.append({"name": name, "mult": mult}) + + def set_lora_from_prompt(self, match): + match = match.split(':') + name = match[0] + + mult = 1.0 + if len(match) == 2: + mult = float(match[1]) + + self.set_lora(name, mult) + def configure_prompt(self, prompt: str) -> str: self.applied_loras = {} - self.prompt = prompt + + for match in re.findall(self.lora_match, prompt): + self.set_lora_from_prompt(match) def found(m): return "" @@ -228,3 +238,4 @@ def __del__(self): for cb in self.hooks: cb.remove() del self.hooks + del self.loras_to_load From 884a5543c77820f3b9fa9bb4d65e3902ab8bca04 Mon Sep 17 00:00:00 2001 From: Jordan Date: Mon, 20 Feb 2023 16:33:53 -0700 Subject: [PATCH 10/30] adjust loader to use a settings dict --- ldm/modules/lora_manager.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index d24316ddc6c..c1719a6db6a 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -120,6 +120,7 @@ def load_lora( class LoraManager: loras: dict[str, LoRA] applied_loras: dict[str, LoRA] + loras_to_load: dict[str, dict] hooks: list[RemovableHandle] def __init__(self, pipe): @@ -132,7 +133,7 @@ def __init__(self, pipe): self.loras = {} self.applied_loras = {} self.hooks = [] - self.loras_to_load = [] + self.loras_to_load = {} def find_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> dict[str, torch.nn.Module]: mapping = {} @@ -153,13 +154,13 @@ def find_modules(prefix, root_module: torch.nn.Module, target_replace_modules) - LORA_PREFIX_UNET, self.unet, UNET_TARGET_REPLACE_MODULE) def _make_hook(self, layer: str): - def hook(module, input, output): + def hook(module, input_h, output): for lora in self.applied_loras.values(): lora_layer = lora.layers.get(layer, None) if lora_layer is None: continue output = output + \ - lora_layer.up(lora_layer.down(*input)) * \ + lora_layer.up(lora_layer.down(*input_h)) * \ lora.multiplier * lora_layer.scale return output return hook @@ -195,19 +196,20 @@ def apply_lora_model(self, name, mult: float = 1.0): self.applied_loras[name] = lora def load_lora(self): - for lora_to_load in self.loras_to_load: - self.apply_lora_model(lora_to_load["name"], lora_to_load["mult"]) + for name, data in self.loras_to_load.items(): + self.apply_lora_model(name, data["mult"]) + + # unload any lora's not defined by loras_to_load + for name in list(self.loras.keys()): + if name not in self.loras_to_load: + self.unload_lora(name) def unload_lora(self, lora_name: str): if lora_name in self.loras: del self.loras[lora_name] def set_lora(self, name, mult: float = 1.0): - if name in self.loras_to_load: - index = self.loras_to_load.index(name) - self.loras_to_load[index]["mult"] = mult - else: - self.loras_to_load.append({"name": name, "mult": mult}) + self.loras_to_load[name] = {"mult": mult} def set_lora_from_prompt(self, match): match = match.split(':') @@ -221,6 +223,7 @@ def set_lora_from_prompt(self, match): def configure_prompt(self, prompt: str) -> str: self.applied_loras = {} + self.loras_to_load = {} for match in re.findall(self.lora_match, prompt): self.set_lora_from_prompt(match) From c3edede73fa38fd21db13a4c175d9ccb8fdea4af Mon Sep 17 00:00:00 2001 From: Jordan Date: Mon, 20 Feb 2023 16:49:59 -0700 Subject: [PATCH 11/30] add notes and adjust functions --- ldm/modules/lora_manager.py | 34 +++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index c1719a6db6a..bb63f501e79 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -98,7 +98,7 @@ def load_lora( value.shape[1], value.shape[0], (1, 1), bias=False) else: print( - f">> Encoundered unknown lora layer module in {name}: {type(value).__name__}") + f">> Encountered unknown lora layer module in {name}: {type(value).__name__}") continue with torch.no_grad(): @@ -125,7 +125,6 @@ class LoraManager: def __init__(self, pipe): self.lora_path = Path(global_models_dir(), 'lora') - self.lora_match = re.compile(r"]+)>") self.unet = pipe.unet self.text_encoder = pipe.text_encoder self.device = torch.device(choose_torch_device()) @@ -208,30 +207,31 @@ def unload_lora(self, lora_name: str): if lora_name in self.loras: del self.loras[lora_name] + # Define a lora to be loaded + # Can be used to define a lora to be loaded outside of prompts def set_lora(self, name, mult: float = 1.0): self.loras_to_load[name] = {"mult": mult} - def set_lora_from_prompt(self, match): - match = match.split(':') - name = match[0] - - mult = 1.0 - if len(match) == 2: - mult = float(match[1]) - - self.set_lora(name, mult) - + # Load the lora from a prompt, syntax is + # Multiplier should be a value between 0.0 and 1.0 def configure_prompt(self, prompt: str) -> str: self.applied_loras = {} self.loras_to_load = {} - for match in re.findall(self.lora_match, prompt): - self.set_lora_from_prompt(match) + lora_match = re.compile(r"]+)>") + + for match in re.findall(lora_match, prompt): + match = match.split(':') + name = match[0] + + mult = 1.0 + if len(match) == 2: + mult = float(match[1]) - def found(m): - return "" + self.set_lora(name, mult) - return re.sub(self.lora_match, found, prompt) + # remove lora and return prompt to avoid the lora prompt causing issues in inference + return re.sub(lora_match, "", prompt) def __del__(self): del self.loras From de890417793bfdd87af7df11b9ea7cbcdc58a1e8 Mon Sep 17 00:00:00 2001 From: Jordan Date: Mon, 20 Feb 2023 17:02:36 -0700 Subject: [PATCH 12/30] optimize functions for unloading --- ldm/modules/lora_manager.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index bb63f501e79..646e8e12373 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -201,7 +201,11 @@ def load_lora(self): # unload any lora's not defined by loras_to_load for name in list(self.loras.keys()): if name not in self.loras_to_load: - self.unload_lora(name) + self.unload_applied_lora(name) + + def unload_applied_lora(self, lora_name: str): + if lora_name in self.applied_loras: + del self.applied_loras[lora_name] def unload_lora(self, lora_name: str): if lora_name in self.loras: @@ -215,8 +219,7 @@ def set_lora(self, name, mult: float = 1.0): # Load the lora from a prompt, syntax is # Multiplier should be a value between 0.0 and 1.0 def configure_prompt(self, prompt: str) -> str: - self.applied_loras = {} - self.loras_to_load = {} + self.clear_loras() lora_match = re.compile(r"]+)>") @@ -233,6 +236,10 @@ def configure_prompt(self, prompt: str) -> str: # remove lora and return prompt to avoid the lora prompt causing issues in inference return re.sub(lora_match, "", prompt) + def clear_loras(self): + self.applied_loras = {} + self.loras_to_load = {} + def __del__(self): del self.loras del self.applied_loras From 3732af63e8cf73c28cab7478d783cd7e347fefbd Mon Sep 17 00:00:00 2001 From: neecapp Date: Mon, 20 Feb 2023 23:06:05 -0500 Subject: [PATCH 13/30] fix prompt --- ldm/generate.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/ldm/generate.py b/ldm/generate.py index c1981d2e23b..a8d13f3a1fb 100644 --- a/ldm/generate.py +++ b/ldm/generate.py @@ -513,6 +513,10 @@ def process_image(image,seed): except AttributeError: pass + # lora MUST process prompt before conditioning + if self.model.lora_manager: + self.model.lora_manager.load_lora() + try: uc, c, extra_conditioning_info = get_uc_and_c_and_ec( prompt, @@ -549,9 +553,6 @@ def process_image(image,seed): else None ) - if self.model.lora_manager: - self.model.lora_manager.load_lora() - results = generator.generate( prompt, iterations=iterations, From e2b6dfeeb9af7dbc43821e3b5cafacf3db1da7fb Mon Sep 17 00:00:00 2001 From: Jordan Date: Mon, 20 Feb 2023 21:33:20 -0700 Subject: [PATCH 14/30] Update generate.py --- ldm/generate.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/ldm/generate.py b/ldm/generate.py index a8d13f3a1fb..8e3b27b2ac2 100644 --- a/ldm/generate.py +++ b/ldm/generate.py @@ -482,6 +482,8 @@ def process_image(image,seed): if self.model.lora_manager: prompt = self.model.lora_manager.configure_prompt(prompt) + # lora MUST process prompt before conditioning + self.model.lora_manager.load_lora() # apply the concepts library to the prompt prompt = self.huggingface_concepts_library.replace_concepts_with_triggers( @@ -513,10 +515,6 @@ def process_image(image,seed): except AttributeError: pass - # lora MUST process prompt before conditioning - if self.model.lora_manager: - self.model.lora_manager.load_lora() - try: uc, c, extra_conditioning_info = get_uc_and_c_and_ec( prompt, From 49c051660229cfc1cb82427a2c2def2f0375b2c1 Mon Sep 17 00:00:00 2001 From: Jordan Date: Mon, 20 Feb 2023 23:45:57 -0700 Subject: [PATCH 15/30] change hook to override --- ldm/modules/lora_manager.py | 126 +++++++++++++++++++++++------------- 1 file changed, 81 insertions(+), 45 deletions(-) diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index 646e8e12373..eaaac19f656 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -15,6 +15,7 @@ class LoRALayer: down: torch.nn.Module def __init__(self, lora_name: str, name: str, rank=4, alpha=1.0): + super().__init__() self.lora_name = lora_name self.name = name self.scale = alpha / rank @@ -37,6 +38,27 @@ def __init__(self, name: str, multiplier=1.0): LORA_PREFIX_TEXT_ENCODER = 'lora_te' +def lora_forward(module, input_h, output): + if len(loaded_loras) == 0: + return output + + lora_name = getattr(module, 'lora_name', None) + for lora in applied_loras.values(): + layer = lora.layers.get(lora_name, None) + if layer is None: + continue + output = output + layer.up(layer.down(*input_h)) * lora.multiplier * layer.scale + return output + + +def lora_linear_forward(self, input_h): + return lora_forward(self, input_h, torch.nn.Linear.forward_before_lora(self, input_h)) + + +def lora_conv2d_forward(self, input_h): + return lora_forward(self, input_h, torch.nn.Conv2d.forward_before_lora(self, input_h)) + + def load_lora( name: str, path_file: Path, @@ -118,10 +140,7 @@ def load_lora( class LoraManager: - loras: dict[str, LoRA] - applied_loras: dict[str, LoRA] - loras_to_load: dict[str, dict] - hooks: list[RemovableHandle] + loras_to_load: dict[str, float] def __init__(self, pipe): self.lora_path = Path(global_models_dir(), 'lora') @@ -129,22 +148,27 @@ def __init__(self, pipe): self.text_encoder = pipe.text_encoder self.device = torch.device(choose_torch_device()) self.dtype = pipe.unet.dtype - self.loras = {} - self.applied_loras = {} - self.hooks = [] self.loras_to_load = {} + if not hasattr(torch.nn.Linear, 'forward_before_lora'): + torch.nn.Linear.forward_before_lora = torch.nn.Linear.forward + + if not hasattr(torch.nn.Conv2d, 'forward_before_lora'): + torch.nn.Conv2d.forward_before_lora = torch.nn.Conv2d.forward + + torch.nn.Linear.forward = lora_linear_forward + torch.nn.Conv2d.forward = lora_conv2d_forward + def find_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> dict[str, torch.nn.Module]: mapping = {} for name, module in root_module.named_modules(): if module.__class__.__name__ in target_replace_modules: for child_name, child_module in module.named_modules(): if child_module.__class__.__name__ == "Linear" or (child_module.__class__.__name__ == "Conv2d" and child_module.kernel_size == (1, 1)): - # It's easier to just convert into "lora" naming instead of trying to inverse map from "lora" -> diffuser lora_name = prefix + '.' + name + '.' + child_name lora_name = lora_name.replace('.', '_') mapping[lora_name] = child_module - self.hooks.append(child_module.register_forward_hook(self._make_hook(lora_name))) + module.lora_name = lora_name return mapping self.text_modules = find_modules( @@ -152,22 +176,10 @@ def find_modules(prefix, root_module: torch.nn.Module, target_replace_modules) - self.unet_modules = find_modules( LORA_PREFIX_UNET, self.unet, UNET_TARGET_REPLACE_MODULE) - def _make_hook(self, layer: str): - def hook(module, input_h, output): - for lora in self.applied_loras.values(): - lora_layer = lora.layers.get(layer, None) - if lora_layer is None: - continue - output = output + \ - lora_layer.up(lora_layer.down(*input_h)) * \ - lora.multiplier * lora_layer.scale - return output - return hook - - def _load_lora(self, name, path_file, multiplier=1.0): + def _load_lora(self, name, path_file, multiplier: float = 1.0): lora = load_lora(name, path_file, self.device, self.dtype, self.text_modules, self.unet_modules, multiplier) - self.loras[name] = lora + loaded_loras[name] = lora return lora def apply_lora_model(self, name, mult: float = 1.0): @@ -175,10 +187,10 @@ def apply_lora_model(self, name, mult: float = 1.0): file = Path(path, "pytorch_lora_weights.bin") if path.is_dir() and file.is_file(): - print(f"loading lora: {path}") - self.unet.load_attn_procs(path.absolute().as_posix()) + print(f"Diffusers lora is currently disabled: {path}") + # print(f"loading lora: {path}") + # self.unet.load_attn_procs(path.absolute().as_posix()) else: - # converting and saving in diffusers format path_file = Path(self.lora_path, f'{name}.ckpt') if Path(self.lora_path, f'{name}.safetensors').exists(): path_file = Path(self.lora_path, f'{name}.safetensors') @@ -187,34 +199,40 @@ def apply_lora_model(self, name, mult: float = 1.0): print(f">> Unable to find lora: {name}") return - lora = self.loras.get(name, None) + lora = loaded_loras.get(name, None) if lora is None: lora = self._load_lora(name, path_file, mult) lora.multiplier = mult - self.applied_loras[name] = lora + applied_loras[name] = lora def load_lora(self): - for name, data in self.loras_to_load.items(): - self.apply_lora_model(name, data["mult"]) + for name, multiplier in self.loras_to_load.items(): + self.apply_lora_model(name, multiplier) # unload any lora's not defined by loras_to_load - for name in list(self.loras.keys()): + for name in list(applied_loras.keys()): if name not in self.loras_to_load: self.unload_applied_lora(name) - def unload_applied_lora(self, lora_name: str): - if lora_name in self.applied_loras: - del self.applied_loras[lora_name] + @staticmethod + def unload_applied_lora(lora_name: str): + if lora_name in applied_loras: + del applied_loras[lora_name] - def unload_lora(self, lora_name: str): - if lora_name in self.loras: - del self.loras[lora_name] + @staticmethod + def unload_lora(lora_name: str): + if lora_name in loaded_loras: + del loaded_loras[lora_name] # Define a lora to be loaded # Can be used to define a lora to be loaded outside of prompts - def set_lora(self, name, mult: float = 1.0): - self.loras_to_load[name] = {"mult": mult} + def set_lora(self, name, multiplier: float = 1.0): + self.loras_to_load[name] = multiplier + + # update the multiplier if the lora was already loaded + if name in loaded_loras: + loaded_loras[name].multiplier = multiplier # Load the lora from a prompt, syntax is # Multiplier should be a value between 0.0 and 1.0 @@ -237,15 +255,33 @@ def configure_prompt(self, prompt: str) -> str: return re.sub(lora_match, "", prompt) def clear_loras(self): - self.applied_loras = {} + clear_applied_loras() self.loras_to_load = {} def __del__(self): - del self.loras - del self.applied_loras + # cleanup overrides + if hasattr(torch.nn.Linear, 'forward_before_lora'): + torch.nn.Linear.forward = torch.nn.Linear.forward_before_lora + del torch.nn.Linear.forward_before_lora + + if hasattr(torch.nn.Conv2d, 'forward_before_lora'): + torch.nn.Conv2d.forward = torch.nn.Conv2d.forward_before_lora + del torch.nn.Conv2d.forward_before_lora + + clear_applied_loras() + clear_loaded_loras() del self.text_modules del self.unet_modules - for cb in self.hooks: - cb.remove() - del self.hooks del self.loras_to_load + + +applied_loras = {} +loaded_loras = {} + + +def clear_applied_loras(): + applied_loras.clear() + + +def clear_loaded_loras(): + loaded_loras.clear() From 5529309e7361339e27d402a4ca64bf2ba5421be3 Mon Sep 17 00:00:00 2001 From: Jordan Date: Tue, 21 Feb 2023 01:34:06 -0700 Subject: [PATCH 16/30] adjusting back to hooks, forcing to be last in execution --- ldm/modules/lora_manager.py | 64 ++++++++++++++++--------------------- 1 file changed, 28 insertions(+), 36 deletions(-) diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index eaaac19f656..0071bf56501 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -38,25 +38,19 @@ def __init__(self, name: str, multiplier=1.0): LORA_PREFIX_TEXT_ENCODER = 'lora_te' -def lora_forward(module, input_h, output): - if len(loaded_loras) == 0: +def lora_forward_hook(name): + def lora_forward(module, input_h, output): + if len(loaded_loras) == 0: + return output + + for lora in applied_loras.values(): + layer = lora.layers.get(name, None) + if layer is None: + continue + output = output + layer.up(layer.down(*input_h)) * lora.multiplier * layer.scale return output - lora_name = getattr(module, 'lora_name', None) - for lora in applied_loras.values(): - layer = lora.layers.get(lora_name, None) - if layer is None: - continue - output = output + layer.up(layer.down(*input_h)) * lora.multiplier * layer.scale - return output - - -def lora_linear_forward(self, input_h): - return lora_forward(self, input_h, torch.nn.Linear.forward_before_lora(self, input_h)) - - -def lora_conv2d_forward(self, input_h): - return lora_forward(self, input_h, torch.nn.Conv2d.forward_before_lora(self, input_h)) + return lora_forward def load_lora( @@ -141,6 +135,7 @@ def load_lora( class LoraManager: loras_to_load: dict[str, float] + hooks: list[RemovableHandle] def __init__(self, pipe): self.lora_path = Path(global_models_dir(), 'lora') @@ -149,26 +144,19 @@ def __init__(self, pipe): self.device = torch.device(choose_torch_device()) self.dtype = pipe.unet.dtype self.loras_to_load = {} - - if not hasattr(torch.nn.Linear, 'forward_before_lora'): - torch.nn.Linear.forward_before_lora = torch.nn.Linear.forward - - if not hasattr(torch.nn.Conv2d, 'forward_before_lora'): - torch.nn.Conv2d.forward_before_lora = torch.nn.Conv2d.forward - - torch.nn.Linear.forward = lora_linear_forward - torch.nn.Conv2d.forward = lora_conv2d_forward + self.hooks = [] def find_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> dict[str, torch.nn.Module]: mapping = {} for name, module in root_module.named_modules(): if module.__class__.__name__ in target_replace_modules: for child_name, child_module in module.named_modules(): - if child_module.__class__.__name__ == "Linear" or (child_module.__class__.__name__ == "Conv2d" and child_module.kernel_size == (1, 1)): + layer_type = child_module.__class__.__name__ + if layer_type == "Linear" or (layer_type == "Conv2d" and child_module.kernel_size == (1, 1)): lora_name = prefix + '.' + name + '.' + child_name lora_name = lora_name.replace('.', '_') mapping[lora_name] = child_module - module.lora_name = lora_name + self.apply_module_forward(child_module, lora_name) return mapping self.text_modules = find_modules( @@ -182,6 +170,12 @@ def _load_lora(self, name, path_file, multiplier: float = 1.0): loaded_loras[name] = lora return lora + def apply_module_forward(self, module, lora_name): + handle = RemovableHandle(module._forward_hooks) + handle.id = 9000 + module._forward_hooks[handle.id] = lora_forward_hook(lora_name) + self.hooks.append(handle) + def apply_lora_model(self, name, mult: float = 1.0): path = Path(self.lora_path, name) file = Path(path, "pytorch_lora_weights.bin") @@ -258,16 +252,14 @@ def clear_loras(self): clear_applied_loras() self.loras_to_load = {} - def __del__(self): - # cleanup overrides - if hasattr(torch.nn.Linear, 'forward_before_lora'): - torch.nn.Linear.forward = torch.nn.Linear.forward_before_lora - del torch.nn.Linear.forward_before_lora + def clear_hooks(self): + for hook in self.hooks: + hook.remove() - if hasattr(torch.nn.Conv2d, 'forward_before_lora'): - torch.nn.Conv2d.forward = torch.nn.Conv2d.forward_before_lora - del torch.nn.Conv2d.forward_before_lora + self.hooks.clear() + def __del__(self): + self.clear_hooks() clear_applied_loras() clear_loaded_loras() del self.text_modules From c669336d6b3958a4befe92e682d72a1e15b8f86d Mon Sep 17 00:00:00 2001 From: Jordan Date: Tue, 21 Feb 2023 02:05:11 -0700 Subject: [PATCH 17/30] Update lora_manager.py --- ldm/modules/lora_manager.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index 0071bf56501..1c4906331f2 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -171,9 +171,7 @@ def _load_lora(self, name, path_file, multiplier: float = 1.0): return lora def apply_module_forward(self, module, lora_name): - handle = RemovableHandle(module._forward_hooks) - handle.id = 9000 - module._forward_hooks[handle.id] = lora_forward_hook(lora_name) + handle = module.register_module_forward_hook(lora_forward_hook(lora_name)) self.hooks.append(handle) def apply_lora_model(self, name, mult: float = 1.0): From 24d92979dbe6ad4205fa0296519f93b615272d36 Mon Sep 17 00:00:00 2001 From: Jordan Date: Tue, 21 Feb 2023 02:08:02 -0700 Subject: [PATCH 18/30] fix typo --- ldm/modules/lora_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index 1c4906331f2..a6c1b948369 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -171,7 +171,7 @@ def _load_lora(self, name, path_file, multiplier: float = 1.0): return lora def apply_module_forward(self, module, lora_name): - handle = module.register_module_forward_hook(lora_forward_hook(lora_name)) + handle = module.register_forward_hook(lora_forward_hook(lora_name)) self.hooks.append(handle) def apply_lora_model(self, name, mult: float = 1.0): From f70b7272f3c14c64d48f2ab700a5afca39666221 Mon Sep 17 00:00:00 2001 From: Jordan Date: Tue, 21 Feb 2023 19:33:39 -0700 Subject: [PATCH 19/30] cleanup / concept of loading through diffusers --- ldm/modules/lora_manager.py | 257 +++++++++++++++++++++++++++--------- 1 file changed, 192 insertions(+), 65 deletions(-) diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index a6c1b948369..b3f85ff0d57 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -5,6 +5,8 @@ from safetensors.torch import load_file import torch from torch.utils.hooks import RemovableHandle +from diffusers.models import UNet2DConditionModel +from transformers import CLIPTextModel class LoRALayer: @@ -15,7 +17,6 @@ class LoRALayer: down: torch.nn.Module def __init__(self, lora_name: str, name: str, rank=4, alpha=1.0): - super().__init__() self.lora_name = lora_name self.name = name self.scale = alpha / rank @@ -24,12 +25,80 @@ def __init__(self, lora_name: str, name: str, rank=4, alpha=1.0): class LoRA: name: str layers: dict[str, LoRALayer] + device: torch.device + dtype: torch.dtype multiplier: float - def __init__(self, name: str, multiplier=1.0): + def __init__(self, name: str, device, dtype, multiplier=1.0): self.name = name self.layers = {} self.multiplier = multiplier + self.device = device + self.dtype = dtype + self.rank = None + self.alpha = None + + def load_from_dict(self, + state_dict, + text_modules: dict[str, torch.nn.Module], + unet_modules: dict[str, torch.nn.Module]): + for key, value in state_dict.items(): + stem, leaf = key.split(".", 1) + + if leaf.endswith("alpha"): + if self.alpha is None: + self.alpha = value.item() + continue + + if stem.startswith(LORA_PREFIX_TEXT_ENCODER): + wrapped = text_modules.get(stem, None) + if wrapped is None: + print(f">> Missing layer: {stem}") + continue + + if self.rank is None and leaf == 'lora_down.weight' and len(value.size()) == 2: + self.rank = value.shape[0] + self.load_lora_layer(stem, leaf, value, wrapped) + continue + elif stem.startswith(LORA_PREFIX_UNET): + wrapped = unet_modules.get(stem, None) + if wrapped is None: + print(f">> Missing layer: {stem}") + continue + + if self.rank is None and leaf == 'lora_down.weight' and len(value.size()) == 2: + self.rank = value.shape[0] + self.load_lora_layer(stem, leaf, value, wrapped) + continue + else: + continue + + def load_lora_layer(self, stem: str, leaf: str, value, wrapped: torch.nn.Module): + layer = self.layers.get(stem, None) + if layer is None: + layer = LoRALayer(self.name, stem, self.rank, self.alpha) + self.layers[stem] = layer + + if type(wrapped) == torch.nn.Linear: + module = torch.nn.Linear(value.shape[1], value.shape[0], bias=False) + elif type(wrapped) == torch.nn.Conv2d: + module = torch.nn.Conv2d(value.shape[1], value.shape[0], (1, 1), bias=False) + else: + print(f">> Encountered unknown lora layer module in {self.name}: {type(value).__name__}") + return + + with torch.no_grad(): + module.weight.copy_(value) + + module.to(device=self.device, dtype=self.dtype) + + if leaf == "lora_up.weight": + layer.up = module + elif leaf == "lora_down.weight": + layer.down = module + else: + print(f">> Encountered unknown layer in lora {self.name}: {leaf}") + return UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"] @@ -37,6 +106,121 @@ def __init__(self, name: str, multiplier=1.0): LORA_PREFIX_UNET = 'lora_unet' LORA_PREFIX_TEXT_ENCODER = 'lora_te' +re_digits = re.compile(r"\d+") +re_unet_transformer_attn_blocks = re.compile( + r"lora_unet_(.+)_blocks_(\d+)_attentions_(\d+)_transformer_blocks_(\d+)_attn(\d+)_(.+).(weight|alpha)" +) +re_unet_mid_blocks = re.compile( + r"lora_unet_mid_block_attentions_(\d+)_(.+).(weight|alpha)" +) +re_unet_transformer_blocks = re.compile( + r"lora_unet_(.+)_blocks_(\d+)_attentions_(\d+)_transformer_blocks_(\d+)_(.+).(weight|alpha)" +) +re_unet_mid_transformer_blocks = re.compile( + r"lora_unet_mid_block_attentions_(\d+)_transformer_blocks_(\d+)_(.+).(weight|alpha)" +) +re_unet_norm_blocks = re.compile( + r"lora_unet_(.+)_blocks_(\d+)_attentions_(\d+)_(.+).(weight|alpha)" +) +re_out = re.compile(r"to_out_(\d+)") +re_processor_weight = re.compile(r"(.+)_(\d+)_(.+)") +re_processor_alpha = re.compile(r"(.+)_(\d+)") + + +def convert_key_to_diffusers(key): + def match(match_list, regex, subject): + r = re.match(regex, subject) + if not r: + return False + + match_list.clear() + match_list.extend([int(x) if re.match(re_digits, x) else x for x in r.groups()]) + return True + + m = [] + + def get_front_block(first, second, third, fourth=None): + if first == "mid": + b_type = f"mid_block" + else: + b_type = f"{first}_blocks.{second}" + + if fourth is None: + return f"{b_type}.attentions.{third}" + + return f"{b_type}.attentions.{third}.transformer_blocks.{fourth}" + + def get_back_block(first, second, third): + second = second.replace(".lora_", "_lora.") + if third == "weight": + bm = [] + if match(bm, re_processor_weight, second): + s_bm = bm[2].split('.') + s_front = f"{bm[0]}_{s_bm[0]}" + s_back = f"{s_bm[1]}" + if int(bm[1]) == 0: + second = f"{s_front}.{s_back}" + else: + second = f"{s_front}.{bm[1]}.{s_back}" + elif third == "alpha": + bma = [] + if match(bma, re_processor_alpha, second): + if int(bma[1]) == 0: + second = f"{bma[0]}" + else: + second = f"{bma[0]}.{bma[1]}" + + if first is None: + return f"processor.{second}.{third}" + + return f"attn{first}.processor.{second}.{third}" + + if match(m, re_unet_transformer_attn_blocks, key): + return f"{get_front_block(m[0], m[1], m[2], m[3])}.{get_back_block(m[4], m[5], m[6])}" + + if match(m, re_unet_transformer_blocks, key): + return f"{get_front_block(m[0], m[1], m[2], m[3])}.{get_back_block(None, m[4], m[5])}" + + if match(m, re_unet_mid_transformer_blocks, key): + return f"{get_front_block('mid', None, m[0], m[1])}.{get_back_block(None, m[2], m[3])}" + + if match(m, re_unet_norm_blocks, key): + return f"{get_front_block(m[0], m[1], m[2])}.{get_back_block(None, m[3], m[4])}" + + if match(m, re_unet_mid_blocks, key): + return f"{get_front_block('mid', None, m[0])}.{get_back_block(None, m[1], m[2])}" + + return key + + +def load_lora_attn( + name: str, + path_file: Path, + unet: UNet2DConditionModel, + text_encoder: CLIPTextModel, + multiplier=1.0 +): + print(f">> Loading lora {name} from {path_file}") + if path_file.suffix == '.safetensors': + checkpoint = load_file(path_file.absolute().as_posix(), device='cpu') + else: + checkpoint = torch.load(path_file, map_location='cpu') + + for key in list(checkpoint.keys()): + if key.startswith(LORA_PREFIX_UNET): + # convert unet keys + checkpoint[convert_key_to_diffusers(key)] = checkpoint.pop(key) + elif key.startswith(LORA_PREFIX_UNET): + # convert text encoder keys (not yet supported) + # state_dict[convert_key_to_diffusers(key)] = state_dict.pop(key) + checkpoint.pop(key) + else: + # remove invalid key + checkpoint.pop(key) + + unet.load_attn_procs(checkpoint) + # text_encoder.load_attn_procs(checkpoint) + def lora_forward_hook(name): def lora_forward(module, input_h, output): @@ -68,67 +252,8 @@ def load_lora( else: checkpoint = torch.load(path_file, map_location='cpu') - lora = LoRA(name, multiplier) - - alpha = None - rank = None - for key, value in checkpoint.items(): - stem, leaf = key.split(".", 1) - - if leaf.endswith("alpha"): - if alpha is None: - alpha = value.item() - continue - - if stem.startswith(LORA_PREFIX_TEXT_ENCODER): - # text encoder layer - wrapped = text_modules.get(stem, None) - if wrapped is None: - print(f">> Missing layer: {stem}") - continue - elif stem.startswith(LORA_PREFIX_UNET): - # unet layer - wrapped = unet_modules.get(stem, None) - if wrapped is None: - print(f">> Missing layer: {stem}") - continue - else: - continue - - if rank is None and leaf == 'lora_down.weight' and len(value.size()) == 2: - rank = value.shape[0] - - if wrapped is None: - continue - - layer = lora.layers.get(stem, None) - if layer is None: - layer = LoRALayer(name, stem, rank, alpha) - lora.layers[stem] = layer - - if type(wrapped) == torch.nn.Linear: - module = torch.nn.Linear( - value.shape[1], value.shape[0], bias=False) - elif type(wrapped) == torch.nn.Conv2d: - module = torch.nn.Conv2d( - value.shape[1], value.shape[0], (1, 1), bias=False) - else: - print( - f">> Encountered unknown lora layer module in {name}: {type(value).__name__}") - continue - - with torch.no_grad(): - module.weight.copy_(value) - - module.to(device=device, dtype=dtype) - - if leaf == "lora_up.weight": - layer.up = module - elif leaf == "lora_down.weight": - layer.down = module - else: - print(f">> Encountered unknown layer in lora {name}: {key}") - continue + lora = LoRA(name, device, dtype, multiplier) + lora.load_from_dict(checkpoint, text_modules, unet_modules) return lora @@ -161,12 +286,14 @@ def find_modules(prefix, root_module: torch.nn.Module, target_replace_modules) - self.text_modules = find_modules( LORA_PREFIX_TEXT_ENCODER, self.text_encoder, TEXT_ENCODER_TARGET_REPLACE_MODULE) + self.unet_modules = find_modules( LORA_PREFIX_UNET, self.unet, UNET_TARGET_REPLACE_MODULE) def _load_lora(self, name, path_file, multiplier: float = 1.0): - lora = load_lora(name, path_file, self.device, self.dtype, - self.text_modules, self.unet_modules, multiplier) + # can be used instead to load through diffusers, once enough support is added + # lora = load_lora_attn(name, path_file, self.unet, self.text_encoder, multiplier) + lora = load_lora(name, path_file, self.device, self.dtype, self.text_modules, self.unet_modules, multiplier) loaded_loras[name] = lora return lora From af3543a8c7ecf91bae24383c43ce287d47427197 Mon Sep 17 00:00:00 2001 From: Jordan Date: Tue, 21 Feb 2023 20:42:40 -0700 Subject: [PATCH 20/30] further cleanup and implement wrapper --- ldm/modules/lora_manager.py | 246 ++++++++++++++++++------------------ 1 file changed, 126 insertions(+), 120 deletions(-) diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index b3f85ff0d57..23d7c9eb223 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -22,26 +22,112 @@ def __init__(self, lora_name: str, name: str, rank=4, alpha=1.0): self.scale = alpha / rank +class LoRAModuleWrapper: + unet: UNet2DConditionModel + text_encoder: CLIPTextModel + + def __init__(self, unet, text_encoder): + self.unet = unet + self.text_encoder = text_encoder + self.hooks = [] + self.text_modules = None + self.unet_modules = None + + self.applied_loras = {} + self.loaded_loras = {} + + self.UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"] + self.TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"] + self.LORA_PREFIX_UNET = 'lora_unet' + self.LORA_PREFIX_TEXT_ENCODER = 'lora_te' + + def find_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> dict[str, torch.nn.Module]: + mapping = {} + for name, module in root_module.named_modules(): + if module.__class__.__name__ in target_replace_modules: + for child_name, child_module in module.named_modules(): + layer_type = child_module.__class__.__name__ + if layer_type == "Linear" or (layer_type == "Conv2d" and child_module.kernel_size == (1, 1)): + lora_name = prefix + '.' + name + '.' + child_name + lora_name = lora_name.replace('.', '_') + mapping[lora_name] = child_module + self.apply_module_forward(child_module, lora_name) + return mapping + + if self.text_modules is None: + self.text_modules = find_modules( + self.LORA_PREFIX_TEXT_ENCODER, + text_encoder, + self.TEXT_ENCODER_TARGET_REPLACE_MODULE + ) + + if self.unet_modules is None: + self.unet_modules = find_modules( + self.LORA_PREFIX_UNET, + unet, + self.UNET_TARGET_REPLACE_MODULE + ) + + def lora_forward_hook(self, name): + wrapper = self + + def lora_forward(module, input_h, output): + if len(wrapper.loaded_loras) == 0: + return output + + for lora in wrapper.applied_loras.values(): + layer = lora.layers.get(name, None) + if layer is None: + continue + output = output + layer.up(layer.down(*input_h)) * lora.multiplier * layer.scale + return output + + return lora_forward + + def apply_module_forward(self, module, name): + handle = module.register_forward_hook(self.lora_forward_hook(name)) + self.hooks.append(handle) + + def clear_hooks(self): + for hook in self.hooks: + hook.remove() + + self.hooks.clear() + + def clear_applied_loras(self): + self.applied_loras.clear() + + def clear_loaded_loras(self): + self.loaded_loras.clear() + + def __del__(self): + self.clear_hooks() + self.clear_applied_loras() + self.clear_loaded_loras() + del self.text_modules + del self.unet_modules + del self.hooks + + class LoRA: name: str layers: dict[str, LoRALayer] device: torch.device dtype: torch.dtype + wrapper: LoRAModuleWrapper multiplier: float - def __init__(self, name: str, device, dtype, multiplier=1.0): + def __init__(self, name: str, device, dtype, wrapper, multiplier=1.0): self.name = name self.layers = {} self.multiplier = multiplier self.device = device self.dtype = dtype + self.wrapper = wrapper self.rank = None self.alpha = None - def load_from_dict(self, - state_dict, - text_modules: dict[str, torch.nn.Module], - unet_modules: dict[str, torch.nn.Module]): + def load_from_dict(self, state_dict): for key, value in state_dict.items(): stem, leaf = key.split(".", 1) @@ -50,8 +136,8 @@ def load_from_dict(self, self.alpha = value.item() continue - if stem.startswith(LORA_PREFIX_TEXT_ENCODER): - wrapped = text_modules.get(stem, None) + if stem.startswith(self.wrapper.LORA_PREFIX_TEXT_ENCODER): + wrapped = self.wrapper.text_modules.get(stem, None) if wrapped is None: print(f">> Missing layer: {stem}") continue @@ -60,8 +146,8 @@ def load_from_dict(self, self.rank = value.shape[0] self.load_lora_layer(stem, leaf, value, wrapped) continue - elif stem.startswith(LORA_PREFIX_UNET): - wrapped = unet_modules.get(stem, None) + elif stem.startswith(self.wrapper.LORA_PREFIX_UNET): + wrapped = self.wrapper.unet_modules.get(stem, None) if wrapped is None: print(f">> Missing layer: {stem}") continue @@ -101,11 +187,6 @@ def load_lora_layer(self, stem: str, leaf: str, value, wrapped: torch.nn.Module) return -UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"] -TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"] -LORA_PREFIX_UNET = 'lora_unet' -LORA_PREFIX_TEXT_ENCODER = 'lora_te' - re_digits = re.compile(r"\d+") re_unet_transformer_attn_blocks = re.compile( r"lora_unet_(.+)_blocks_(\d+)_attentions_(\d+)_transformer_blocks_(\d+)_attn(\d+)_(.+).(weight|alpha)" @@ -196,8 +277,7 @@ def get_back_block(first, second, third): def load_lora_attn( name: str, path_file: Path, - unet: UNet2DConditionModel, - text_encoder: CLIPTextModel, + wrapper: LoRAModuleWrapper, multiplier=1.0 ): print(f">> Loading lora {name} from {path_file}") @@ -207,10 +287,10 @@ def load_lora_attn( checkpoint = torch.load(path_file, map_location='cpu') for key in list(checkpoint.keys()): - if key.startswith(LORA_PREFIX_UNET): + if key.startswith(wrapper.LORA_PREFIX_UNET): # convert unet keys checkpoint[convert_key_to_diffusers(key)] = checkpoint.pop(key) - elif key.startswith(LORA_PREFIX_UNET): + elif key.startswith(wrapper.LORA_PREFIX_UNET): # convert text encoder keys (not yet supported) # state_dict[convert_key_to_diffusers(key)] = state_dict.pop(key) checkpoint.pop(key) @@ -218,44 +298,8 @@ def load_lora_attn( # remove invalid key checkpoint.pop(key) - unet.load_attn_procs(checkpoint) - # text_encoder.load_attn_procs(checkpoint) - - -def lora_forward_hook(name): - def lora_forward(module, input_h, output): - if len(loaded_loras) == 0: - return output - - for lora in applied_loras.values(): - layer = lora.layers.get(name, None) - if layer is None: - continue - output = output + layer.up(layer.down(*input_h)) * lora.multiplier * layer.scale - return output - - return lora_forward - - -def load_lora( - name: str, - path_file: Path, - device: torch.device, - dtype: torch.dtype, - text_modules: dict[str, torch.nn.Module], - unet_modules: dict[str, torch.nn.Module], - multiplier=1.0 -): - print(f">> Loading lora {name} from {path_file}") - if path_file.suffix == '.safetensors': - checkpoint = load_file(path_file.absolute().as_posix(), device='cpu') - else: - checkpoint = torch.load(path_file, map_location='cpu') - - lora = LoRA(name, device, dtype, multiplier) - lora.load_from_dict(checkpoint, text_modules, unet_modules) - - return lora + wrapper.unet.load_attn_procs(checkpoint) + # wrapper.text_encoder.load_attn_procs(checkpoint) class LoraManager: @@ -269,38 +313,24 @@ def __init__(self, pipe): self.device = torch.device(choose_torch_device()) self.dtype = pipe.unet.dtype self.loras_to_load = {} - self.hooks = [] + self.wrapper = LoRAModuleWrapper(pipe.unet, pipe.text_encoder) - def find_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> dict[str, torch.nn.Module]: - mapping = {} - for name, module in root_module.named_modules(): - if module.__class__.__name__ in target_replace_modules: - for child_name, child_module in module.named_modules(): - layer_type = child_module.__class__.__name__ - if layer_type == "Linear" or (layer_type == "Conv2d" and child_module.kernel_size == (1, 1)): - lora_name = prefix + '.' + name + '.' + child_name - lora_name = lora_name.replace('.', '_') - mapping[lora_name] = child_module - self.apply_module_forward(child_module, lora_name) - return mapping + def load_lora_module(self, name, path_file, multiplier: float = 1.0): + # can be used instead to load through diffusers, once enough support is added + # lora = load_lora_attn(name, path_file, self.wrapper, multiplier) - self.text_modules = find_modules( - LORA_PREFIX_TEXT_ENCODER, self.text_encoder, TEXT_ENCODER_TARGET_REPLACE_MODULE) + print(f">> Loading lora {name} from {path_file}") + if path_file.suffix == '.safetensors': + checkpoint = load_file(path_file.absolute().as_posix(), device='cpu') + else: + checkpoint = torch.load(path_file, map_location='cpu') - self.unet_modules = find_modules( - LORA_PREFIX_UNET, self.unet, UNET_TARGET_REPLACE_MODULE) + lora = LoRA(name, self.device, self.dtype, self.wrapper, multiplier) + lora.load_from_dict(checkpoint) + self.wrapper.loaded_loras[name] = lora - def _load_lora(self, name, path_file, multiplier: float = 1.0): - # can be used instead to load through diffusers, once enough support is added - # lora = load_lora_attn(name, path_file, self.unet, self.text_encoder, multiplier) - lora = load_lora(name, path_file, self.device, self.dtype, self.text_modules, self.unet_modules, multiplier) - loaded_loras[name] = lora return lora - def apply_module_forward(self, module, lora_name): - handle = module.register_forward_hook(lora_forward_hook(lora_name)) - self.hooks.append(handle) - def apply_lora_model(self, name, mult: float = 1.0): path = Path(self.lora_path, name) file = Path(path, "pytorch_lora_weights.bin") @@ -318,31 +348,30 @@ def apply_lora_model(self, name, mult: float = 1.0): print(f">> Unable to find lora: {name}") return - lora = loaded_loras.get(name, None) + lora = self.wrapper.loaded_loras.get(name, None) if lora is None: - lora = self._load_lora(name, path_file, mult) + lora = self.load_lora_module(name, path_file, mult) lora.multiplier = mult - applied_loras[name] = lora + self.wrapper.applied_loras[name] = lora def load_lora(self): + print(self.loras_to_load) for name, multiplier in self.loras_to_load.items(): self.apply_lora_model(name, multiplier) # unload any lora's not defined by loras_to_load - for name in list(applied_loras.keys()): + for name in list(self.wrapper.applied_loras.keys()): if name not in self.loras_to_load: self.unload_applied_lora(name) - @staticmethod - def unload_applied_lora(lora_name: str): - if lora_name in applied_loras: - del applied_loras[lora_name] + def unload_applied_lora(self, lora_name: str): + if lora_name in self.wrapper.applied_loras: + del self.wrapper.applied_loras[lora_name] - @staticmethod - def unload_lora(lora_name: str): - if lora_name in loaded_loras: - del loaded_loras[lora_name] + def unload_lora(self, lora_name: str): + if lora_name in self.wrapper.loaded_loras: + del self.wrapper.loaded_loras[lora_name] # Define a lora to be loaded # Can be used to define a lora to be loaded outside of prompts @@ -350,8 +379,8 @@ def set_lora(self, name, multiplier: float = 1.0): self.loras_to_load[name] = multiplier # update the multiplier if the lora was already loaded - if name in loaded_loras: - loaded_loras[name].multiplier = multiplier + if name in self.wrapper.loaded_loras: + self.wrapper.loaded_loras[name].multiplier = multiplier # Load the lora from a prompt, syntax is # Multiplier should be a value between 0.0 and 1.0 @@ -374,31 +403,8 @@ def configure_prompt(self, prompt: str) -> str: return re.sub(lora_match, "", prompt) def clear_loras(self): - clear_applied_loras() + self.wrapper.clear_applied_loras() self.loras_to_load = {} - def clear_hooks(self): - for hook in self.hooks: - hook.remove() - - self.hooks.clear() - def __del__(self): - self.clear_hooks() - clear_applied_loras() - clear_loaded_loras() - del self.text_modules - del self.unet_modules del self.loras_to_load - - -applied_loras = {} -loaded_loras = {} - - -def clear_applied_loras(): - applied_loras.clear() - - -def clear_loaded_loras(): - loaded_loras.clear() From cd333e414bc88fed16b17f018188b07f76eb0e59 Mon Sep 17 00:00:00 2001 From: Jordan Date: Tue, 21 Feb 2023 21:38:15 -0700 Subject: [PATCH 21/30] move key converter to wrapper --- ldm/modules/lora_manager.py | 179 ++++++++++++++++++------------------ 1 file changed, 88 insertions(+), 91 deletions(-) diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index 23d7c9eb223..b62feae2fb5 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -25,6 +25,7 @@ def __init__(self, lora_name: str, name: str, rank=4, alpha=1.0): class LoRAModuleWrapper: unet: UNet2DConditionModel text_encoder: CLIPTextModel + hooks: list[RemovableHandle] def __init__(self, unet, text_encoder): self.unet = unet @@ -41,6 +42,26 @@ def __init__(self, unet, text_encoder): self.LORA_PREFIX_UNET = 'lora_unet' self.LORA_PREFIX_TEXT_ENCODER = 'lora_te' + self.re_digits = re.compile(r"\d+") + self.re_unet_transformer_attn_blocks = re.compile( + r"lora_unet_(.+)_blocks_(\d+)_attentions_(\d+)_transformer_blocks_(\d+)_attn(\d+)_(.+).(weight|alpha)" + ) + self.re_unet_mid_blocks = re.compile( + r"lora_unet_mid_block_attentions_(\d+)_(.+).(weight|alpha)" + ) + self.re_unet_transformer_blocks = re.compile( + r"lora_unet_(.+)_blocks_(\d+)_attentions_(\d+)_transformer_blocks_(\d+)_(.+).(weight|alpha)" + ) + self.re_unet_mid_transformer_blocks = re.compile( + r"lora_unet_mid_block_attentions_(\d+)_transformer_blocks_(\d+)_(.+).(weight|alpha)" + ) + self.re_unet_norm_blocks = re.compile( + r"lora_unet_(.+)_blocks_(\d+)_attentions_(\d+)_(.+).(weight|alpha)" + ) + self.re_out = re.compile(r"to_out_(\d+)") + self.re_processor_weight = re.compile(r"(.+)_(\d+)_(.+)") + self.re_processor_alpha = re.compile(r"(.+)_(\d+)") + def find_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> dict[str, torch.nn.Module]: mapping = {} for name, module in root_module.named_modules(): @@ -68,6 +89,71 @@ def find_modules(prefix, root_module: torch.nn.Module, target_replace_modules) - self.UNET_TARGET_REPLACE_MODULE ) + def convert_key_to_diffusers(self, key): + def match(match_list, regex, subject): + r = re.match(regex, subject) + if not r: + return False + + match_list.clear() + match_list.extend([int(x) if re.match(self.re_digits, x) else x for x in r.groups()]) + return True + + m = [] + + def get_front_block(first, second, third, fourth=None): + if first == "mid": + b_type = f"mid_block" + else: + b_type = f"{first}_blocks.{second}" + + if fourth is None: + return f"{b_type}.attentions.{third}" + + return f"{b_type}.attentions.{third}.transformer_blocks.{fourth}" + + def get_back_block(first, second, third): + second = second.replace(".lora_", "_lora.") + if third == "weight": + bm = [] + if match(bm, self.re_processor_weight, second): + s_bm = bm[2].split('.') + s_front = f"{bm[0]}_{s_bm[0]}" + s_back = f"{s_bm[1]}" + if int(bm[1]) == 0: + second = f"{s_front}.{s_back}" + else: + second = f"{s_front}.{bm[1]}.{s_back}" + elif third == "alpha": + bma = [] + if match(bma, self.re_processor_alpha, second): + if int(bma[1]) == 0: + second = f"{bma[0]}" + else: + second = f"{bma[0]}.{bma[1]}" + + if first is None: + return f"processor.{second}.{third}" + + return f"attn{first}.processor.{second}.{third}" + + if match(m, self.re_unet_transformer_attn_blocks, key): + return f"{get_front_block(m[0], m[1], m[2], m[3])}.{get_back_block(m[4], m[5], m[6])}" + + if match(m, self.re_unet_transformer_blocks, key): + return f"{get_front_block(m[0], m[1], m[2], m[3])}.{get_back_block(None, m[4], m[5])}" + + if match(m, self.re_unet_mid_transformer_blocks, key): + return f"{get_front_block('mid', None, m[0], m[1])}.{get_back_block(None, m[2], m[3])}" + + if match(m, self.re_unet_norm_blocks, key): + return f"{get_front_block(m[0], m[1], m[2])}.{get_back_block(None, m[3], m[4])}" + + if match(m, self.re_unet_mid_blocks, key): + return f"{get_front_block('mid', None, m[0])}.{get_back_block(None, m[1], m[2])}" + + return key + def lora_forward_hook(self, name): wrapper = self @@ -187,93 +273,6 @@ def load_lora_layer(self, stem: str, leaf: str, value, wrapped: torch.nn.Module) return -re_digits = re.compile(r"\d+") -re_unet_transformer_attn_blocks = re.compile( - r"lora_unet_(.+)_blocks_(\d+)_attentions_(\d+)_transformer_blocks_(\d+)_attn(\d+)_(.+).(weight|alpha)" -) -re_unet_mid_blocks = re.compile( - r"lora_unet_mid_block_attentions_(\d+)_(.+).(weight|alpha)" -) -re_unet_transformer_blocks = re.compile( - r"lora_unet_(.+)_blocks_(\d+)_attentions_(\d+)_transformer_blocks_(\d+)_(.+).(weight|alpha)" -) -re_unet_mid_transformer_blocks = re.compile( - r"lora_unet_mid_block_attentions_(\d+)_transformer_blocks_(\d+)_(.+).(weight|alpha)" -) -re_unet_norm_blocks = re.compile( - r"lora_unet_(.+)_blocks_(\d+)_attentions_(\d+)_(.+).(weight|alpha)" -) -re_out = re.compile(r"to_out_(\d+)") -re_processor_weight = re.compile(r"(.+)_(\d+)_(.+)") -re_processor_alpha = re.compile(r"(.+)_(\d+)") - - -def convert_key_to_diffusers(key): - def match(match_list, regex, subject): - r = re.match(regex, subject) - if not r: - return False - - match_list.clear() - match_list.extend([int(x) if re.match(re_digits, x) else x for x in r.groups()]) - return True - - m = [] - - def get_front_block(first, second, third, fourth=None): - if first == "mid": - b_type = f"mid_block" - else: - b_type = f"{first}_blocks.{second}" - - if fourth is None: - return f"{b_type}.attentions.{third}" - - return f"{b_type}.attentions.{third}.transformer_blocks.{fourth}" - - def get_back_block(first, second, third): - second = second.replace(".lora_", "_lora.") - if third == "weight": - bm = [] - if match(bm, re_processor_weight, second): - s_bm = bm[2].split('.') - s_front = f"{bm[0]}_{s_bm[0]}" - s_back = f"{s_bm[1]}" - if int(bm[1]) == 0: - second = f"{s_front}.{s_back}" - else: - second = f"{s_front}.{bm[1]}.{s_back}" - elif third == "alpha": - bma = [] - if match(bma, re_processor_alpha, second): - if int(bma[1]) == 0: - second = f"{bma[0]}" - else: - second = f"{bma[0]}.{bma[1]}" - - if first is None: - return f"processor.{second}.{third}" - - return f"attn{first}.processor.{second}.{third}" - - if match(m, re_unet_transformer_attn_blocks, key): - return f"{get_front_block(m[0], m[1], m[2], m[3])}.{get_back_block(m[4], m[5], m[6])}" - - if match(m, re_unet_transformer_blocks, key): - return f"{get_front_block(m[0], m[1], m[2], m[3])}.{get_back_block(None, m[4], m[5])}" - - if match(m, re_unet_mid_transformer_blocks, key): - return f"{get_front_block('mid', None, m[0], m[1])}.{get_back_block(None, m[2], m[3])}" - - if match(m, re_unet_norm_blocks, key): - return f"{get_front_block(m[0], m[1], m[2])}.{get_back_block(None, m[3], m[4])}" - - if match(m, re_unet_mid_blocks, key): - return f"{get_front_block('mid', None, m[0])}.{get_back_block(None, m[1], m[2])}" - - return key - - def load_lora_attn( name: str, path_file: Path, @@ -289,10 +288,10 @@ def load_lora_attn( for key in list(checkpoint.keys()): if key.startswith(wrapper.LORA_PREFIX_UNET): # convert unet keys - checkpoint[convert_key_to_diffusers(key)] = checkpoint.pop(key) + checkpoint[wrapper.convert_key_to_diffusers(key)] = checkpoint.pop(key) elif key.startswith(wrapper.LORA_PREFIX_UNET): # convert text encoder keys (not yet supported) - # state_dict[convert_key_to_diffusers(key)] = state_dict.pop(key) + # state_dict[wrapper.convert_key_to_diffusers(key)] = state_dict.pop(key) checkpoint.pop(key) else: # remove invalid key @@ -304,7 +303,6 @@ def load_lora_attn( class LoraManager: loras_to_load: dict[str, float] - hooks: list[RemovableHandle] def __init__(self, pipe): self.lora_path = Path(global_models_dir(), 'lora') @@ -356,7 +354,6 @@ def apply_lora_model(self, name, mult: float = 1.0): self.wrapper.applied_loras[name] = lora def load_lora(self): - print(self.loras_to_load) for name, multiplier in self.loras_to_load.items(): self.apply_lora_model(name, multiplier) From 71972c370922e8f7e16bcc9e0fb6ff7d64368c14 Mon Sep 17 00:00:00 2001 From: Jordan Date: Thu, 23 Feb 2023 01:44:13 -0700 Subject: [PATCH 22/30] re-enable load attn procs support (no multiplier) --- ldm/modules/lora_manager.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index b62feae2fb5..e9dfe688ecd 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -334,9 +334,8 @@ def apply_lora_model(self, name, mult: float = 1.0): file = Path(path, "pytorch_lora_weights.bin") if path.is_dir() and file.is_file(): - print(f"Diffusers lora is currently disabled: {path}") - # print(f"loading lora: {path}") - # self.unet.load_attn_procs(path.absolute().as_posix()) + print(f"loading lora: {path}") + self.unet.load_attn_procs(path.absolute().as_posix()) else: path_file = Path(self.lora_path, f'{name}.ckpt') if Path(self.lora_path, f'{name}.safetensors').exists(): From f64a4db5fa2c34f842cbd42061662d710230a03c Mon Sep 17 00:00:00 2001 From: Jordan Date: Thu, 23 Feb 2023 05:56:39 -0700 Subject: [PATCH 23/30] setup legacy class to abstract hacky logic for none diffusers lora and format prompt for compel --- ldm/modules/lora_manager.py | 127 +++++++++++++++++------------------- 1 file changed, 60 insertions(+), 67 deletions(-) diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index e9dfe688ecd..902a2cb5781 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -273,45 +273,14 @@ def load_lora_layer(self, stem: str, leaf: str, value, wrapped: torch.nn.Module) return -def load_lora_attn( - name: str, - path_file: Path, - wrapper: LoRAModuleWrapper, - multiplier=1.0 -): - print(f">> Loading lora {name} from {path_file}") - if path_file.suffix == '.safetensors': - checkpoint = load_file(path_file.absolute().as_posix(), device='cpu') - else: - checkpoint = torch.load(path_file, map_location='cpu') - - for key in list(checkpoint.keys()): - if key.startswith(wrapper.LORA_PREFIX_UNET): - # convert unet keys - checkpoint[wrapper.convert_key_to_diffusers(key)] = checkpoint.pop(key) - elif key.startswith(wrapper.LORA_PREFIX_UNET): - # convert text encoder keys (not yet supported) - # state_dict[wrapper.convert_key_to_diffusers(key)] = state_dict.pop(key) - checkpoint.pop(key) - else: - # remove invalid key - checkpoint.pop(key) - - wrapper.unet.load_attn_procs(checkpoint) - # wrapper.text_encoder.load_attn_procs(checkpoint) - - -class LoraManager: - loras_to_load: dict[str, float] - - def __init__(self, pipe): - self.lora_path = Path(global_models_dir(), 'lora') +class LegacyLora: + def __init__(self, pipe, lora_path): self.unet = pipe.unet + self.lora_path = lora_path + self.wrapper = LoRAModuleWrapper(pipe.unet, pipe.text_encoder) self.text_encoder = pipe.text_encoder self.device = torch.device(choose_torch_device()) self.dtype = pipe.unet.dtype - self.loras_to_load = {} - self.wrapper = LoRAModuleWrapper(pipe.unet, pipe.text_encoder) def load_lora_module(self, name, path_file, multiplier: float = 1.0): # can be used instead to load through diffusers, once enough support is added @@ -330,35 +299,25 @@ def load_lora_module(self, name, path_file, multiplier: float = 1.0): return lora def apply_lora_model(self, name, mult: float = 1.0): - path = Path(self.lora_path, name) - file = Path(path, "pytorch_lora_weights.bin") - - if path.is_dir() and file.is_file(): - print(f"loading lora: {path}") - self.unet.load_attn_procs(path.absolute().as_posix()) - else: - path_file = Path(self.lora_path, f'{name}.ckpt') - if Path(self.lora_path, f'{name}.safetensors').exists(): - path_file = Path(self.lora_path, f'{name}.safetensors') + path_file = Path(self.lora_path, f'{name}.ckpt') + if Path(self.lora_path, f'{name}.safetensors').exists(): + path_file = Path(self.lora_path, f'{name}.safetensors') - if not path_file.exists(): - print(f">> Unable to find lora: {name}") - return + if not path_file.exists(): + print(f">> Unable to find lora: {name}") + return - lora = self.wrapper.loaded_loras.get(name, None) - if lora is None: - lora = self.load_lora_module(name, path_file, mult) + lora = self.wrapper.loaded_loras.get(name, None) + if lora is None: + lora = self.load_lora_module(name, path_file, mult) - lora.multiplier = mult - self.wrapper.applied_loras[name] = lora - - def load_lora(self): - for name, multiplier in self.loras_to_load.items(): - self.apply_lora_model(name, multiplier) + lora.multiplier = mult + self.wrapper.applied_loras[name] = lora + def unload_applied_loras(self, loras_to_load): # unload any lora's not defined by loras_to_load for name in list(self.wrapper.applied_loras.keys()): - if name not in self.loras_to_load: + if name not in loras_to_load: self.unload_applied_lora(name) def unload_applied_lora(self, lora_name: str): @@ -369,29 +328,63 @@ def unload_lora(self, lora_name: str): if lora_name in self.wrapper.loaded_loras: del self.wrapper.loaded_loras[lora_name] - # Define a lora to be loaded - # Can be used to define a lora to be loaded outside of prompts def set_lora(self, name, multiplier: float = 1.0): - self.loras_to_load[name] = multiplier - # update the multiplier if the lora was already loaded if name in self.wrapper.loaded_loras: self.wrapper.loaded_loras[name].multiplier = multiplier + def clear_loras(self): + self.wrapper.clear_applied_loras() + + +class LoraManager: + loras_to_load: dict[str, float] + + def __init__(self, pipe): + self.lora_path = Path(global_models_dir(), 'lora') + self.unet = pipe.unet + self.loras_to_load = {} + # Legacy class handles lora not generated through diffusers + self.legacy = LegacyLora(pipe, self.lora_path) + + def apply_lora_model(self, name, mult: float = 1.0): + path = Path(self.lora_path, name) + file = Path(path, "pytorch_lora_weights.bin") + + if path.is_dir() and file.is_file(): + print(f"loading lora: {path}") + self.unet.load_attn_procs(path.absolute().as_posix()) + else: + self.legacy.apply_lora_model(name, mult) + + def load_lora(self): + for name, multiplier in self.loras_to_load.items(): + self.apply_lora_model(name, multiplier) + + self.legacy.unload_applied_loras(self.loras_to_load) + + # Define a lora to be loaded + # Can be used to define a lora to be loaded outside of prompts + def set_lora(self, name, multiplier: float = 1.0): + self.loras_to_load[name] = multiplier + self.legacy.set_lora(name, multiplier) + # Load the lora from a prompt, syntax is # Multiplier should be a value between 0.0 and 1.0 def configure_prompt(self, prompt: str) -> str: self.clear_loras() - lora_match = re.compile(r"]+)>") + # lora_match = re.compile(r"]+)>") + lora_match = re.compile(r"withLora\(([a-zA-Z\,\d]+)\)") for match in re.findall(lora_match, prompt): - match = match.split(':') - name = match[0] + # match = match.split(':') + match = match.split(',') + name = match[0].strip() mult = 1.0 if len(match) == 2: - mult = float(match[1]) + mult = float(match[1].strip()) self.set_lora(name, mult) @@ -399,8 +392,8 @@ def configure_prompt(self, prompt: str) -> str: return re.sub(lora_match, "", prompt) def clear_loras(self): - self.wrapper.clear_applied_loras() self.loras_to_load = {} + self.legacy.clear_loras() def __del__(self): del self.loras_to_load From 6a1129ab64c66429f7cca2b5d84ca4530cb0fb57 Mon Sep 17 00:00:00 2001 From: Jordan Date: Thu, 23 Feb 2023 16:48:33 -0700 Subject: [PATCH 24/30] switch all none diffusers stuff to legacy, and load through compel prompts --- ldm/generate.py | 5 +- ldm/invoke/conditioning.py | 2 + ldm/modules/lora_manager.py | 91 ++++++++++++++++++------------------- 3 files changed, 50 insertions(+), 48 deletions(-) diff --git a/ldm/generate.py b/ldm/generate.py index ba6dec9d5bc..9a7badb42e1 100644 --- a/ldm/generate.py +++ b/ldm/generate.py @@ -488,10 +488,11 @@ def process_image(image,seed): self.sampler_name = sampler_name self._set_sampler() + # To try and load lora not trained through diffusers if self.model.lora_manager: - prompt = self.model.lora_manager.configure_prompt(prompt) + prompt = self.model.lora_manager.configure_prompt_legacy(prompt) # lora MUST process prompt before conditioning - self.model.lora_manager.load_lora() + self.model.lora_manager.load_lora_legacy() # apply the concepts library to the prompt prompt = self.huggingface_concepts_library.replace_concepts_with_triggers( diff --git a/ldm/invoke/conditioning.py b/ldm/invoke/conditioning.py index db50ce40763..94e90cab559 100644 --- a/ldm/invoke/conditioning.py +++ b/ldm/invoke/conditioning.py @@ -59,6 +59,8 @@ def get_uc_and_c_and_ec(prompt_string, model, log_tokens=False, skip_normalize_l positive_prompt = legacy_blend else: positive_prompt = Compel.parse_prompt_string(positive_prompt_string) + if model.lora_manager: + model.lora_manager.load_lora_compel(positive_prompt.lora_weights) negative_prompt: FlattenedPrompt|Blend = Compel.parse_prompt_string(negative_prompt_string) if log_tokens or getattr(Globals, "log_tokenization", False): diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index 902a2cb5781..5c4d52d01f2 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -273,7 +273,7 @@ def load_lora_layer(self, stem: str, leaf: str, value, wrapped: torch.nn.Module) return -class LegacyLora: +class LegacyLoraManager: def __init__(self, pipe, lora_path): self.unet = pipe.unet self.lora_path = lora_path @@ -281,6 +281,7 @@ def __init__(self, pipe, lora_path): self.text_encoder = pipe.text_encoder self.device = torch.device(choose_torch_device()) self.dtype = pipe.unet.dtype + self.loras_to_load = {} def load_lora_module(self, name, path_file, multiplier: float = 1.0): # can be used instead to load through diffusers, once enough support is added @@ -298,6 +299,26 @@ def load_lora_module(self, name, path_file, multiplier: float = 1.0): return lora + def configure_prompt(self, prompt: str) -> str: + self.clear_loras() + + # lora_match = re.compile(r"]+)>") + lora_match = re.compile(r"withLoraLegacy\(([a-zA-Z\,\d]+)\)") + + for match in re.findall(lora_match, prompt): + # match = match.split(':') + match = match.split(',') + name = match[0].strip() + + mult = 1.0 + if len(match) == 2: + mult = float(match[1].strip()) + + self.set_lora(name, mult) + + # remove lora and return prompt to avoid the lora prompt causing issues in inference + return re.sub(lora_match, "", prompt) + def apply_lora_model(self, name, mult: float = 1.0): path_file = Path(self.lora_path, f'{name}.ckpt') if Path(self.lora_path, f'{name}.safetensors').exists(): @@ -314,6 +335,10 @@ def apply_lora_model(self, name, mult: float = 1.0): lora.multiplier = mult self.wrapper.applied_loras[name] = lora + def load_lora(self): + for name, multiplier in self.loras_to_load.items(): + self.apply_lora_model(name, multiplier) + def unload_applied_loras(self, loras_to_load): # unload any lora's not defined by loras_to_load for name in list(self.wrapper.applied_loras.keys()): @@ -329,71 +354,45 @@ def unload_lora(self, lora_name: str): del self.wrapper.loaded_loras[lora_name] def set_lora(self, name, multiplier: float = 1.0): + self.loras_to_load[name] = multiplier + # update the multiplier if the lora was already loaded if name in self.wrapper.loaded_loras: self.wrapper.loaded_loras[name].multiplier = multiplier def clear_loras(self): + self.loras_to_load = {} self.wrapper.clear_applied_loras() + def __del__(self): + del self.loras_to_load + class LoraManager: - loras_to_load: dict[str, float] - def __init__(self, pipe): self.lora_path = Path(global_models_dir(), 'lora') self.unet = pipe.unet - self.loras_to_load = {} # Legacy class handles lora not generated through diffusers - self.legacy = LegacyLora(pipe, self.lora_path) + self.legacy = LegacyLoraManager(pipe, self.lora_path) - def apply_lora_model(self, name, mult: float = 1.0): + def apply_lora_model(self, name): path = Path(self.lora_path, name) file = Path(path, "pytorch_lora_weights.bin") if path.is_dir() and file.is_file(): - print(f"loading lora: {path}") + print(f">> Loading LoRA: {path}") self.unet.load_attn_procs(path.absolute().as_posix()) else: - self.legacy.apply_lora_model(name, mult) - - def load_lora(self): - for name, multiplier in self.loras_to_load.items(): - self.apply_lora_model(name, multiplier) - - self.legacy.unload_applied_loras(self.loras_to_load) - - # Define a lora to be loaded - # Can be used to define a lora to be loaded outside of prompts - def set_lora(self, name, multiplier: float = 1.0): - self.loras_to_load[name] = multiplier - self.legacy.set_lora(name, multiplier) + print(f">> Unable to find valid LoRA at: {path}") - # Load the lora from a prompt, syntax is - # Multiplier should be a value between 0.0 and 1.0 - def configure_prompt(self, prompt: str) -> str: - self.clear_loras() + def load_lora_compel(self, lora_weights: list): + if len(lora_weights) > 0: + for lora in lora_weights: + self.apply_lora_model(lora.model) - # lora_match = re.compile(r"]+)>") - lora_match = re.compile(r"withLora\(([a-zA-Z\,\d]+)\)") + # Legacy functions, to pipe to LoraLegacyManager + def configure_prompt_legacy(self, prompt: str) -> str: + return self.legacy.configure_prompt(prompt) - for match in re.findall(lora_match, prompt): - # match = match.split(':') - match = match.split(',') - name = match[0].strip() - - mult = 1.0 - if len(match) == 2: - mult = float(match[1].strip()) - - self.set_lora(name, mult) - - # remove lora and return prompt to avoid the lora prompt causing issues in inference - return re.sub(lora_match, "", prompt) - - def clear_loras(self): - self.loras_to_load = {} - self.legacy.clear_loras() - - def __del__(self): - del self.loras_to_load + def load_lora_legacy(self): + self.legacy.load_lora() From b69f9d4af18ba02cae55d6f3b771f63b0c62795f Mon Sep 17 00:00:00 2001 From: Jordan Date: Thu, 23 Feb 2023 17:30:34 -0700 Subject: [PATCH 25/30] initial setup of cross attention --- ldm/invoke/conditioning.py | 2 +- ldm/invoke/generator/diffusers_pipeline.py | 7 +++++-- ldm/models/diffusion/shared_invokeai_diffusion.py | 7 ++++++- ldm/modules/lora_manager.py | 15 +++++++++++++-- 4 files changed, 25 insertions(+), 6 deletions(-) diff --git a/ldm/invoke/conditioning.py b/ldm/invoke/conditioning.py index 94e90cab559..32b19fab655 100644 --- a/ldm/invoke/conditioning.py +++ b/ldm/invoke/conditioning.py @@ -60,7 +60,7 @@ def get_uc_and_c_and_ec(prompt_string, model, log_tokens=False, skip_normalize_l else: positive_prompt = Compel.parse_prompt_string(positive_prompt_string) if model.lora_manager: - model.lora_manager.load_lora_compel(positive_prompt.lora_weights) + model.lora_manager.set_loras_compel(positive_prompt.lora_weights) negative_prompt: FlattenedPrompt|Blend = Compel.parse_prompt_string(negative_prompt_string) if log_tokens or getattr(Globals, "log_tokenization", False): diff --git a/ldm/invoke/generator/diffusers_pipeline.py b/ldm/invoke/generator/diffusers_pipeline.py index a34a503dfc4..7fd4b5880ed 100644 --- a/ldm/invoke/generator/diffusers_pipeline.py +++ b/ldm/invoke/generator/diffusers_pipeline.py @@ -290,12 +290,15 @@ def __init__( safety_checker=safety_checker, feature_extractor=feature_extractor, ) - self.invokeai_diffuser = InvokeAIDiffuserComponent(self.unet, self._unet_forward, is_running_diffusers=True) + self.lora_manager = LoraManager(self) + self.invokeai_diffuser = InvokeAIDiffuserComponent(self.unet, + self._unet_forward, + self.lora_manager, + is_running_diffusers=True) use_full_precision = (precision == 'float32' or precision == 'autocast') self.textual_inversion_manager = TextualInversionManager(tokenizer=self.tokenizer, text_encoder=self.text_encoder, full_precision=use_full_precision) - self.lora_manager = LoraManager(self) # InvokeAI's interface for text embeddings and whatnot self.embeddings_provider = EmbeddingsProvider( diff --git a/ldm/models/diffusion/shared_invokeai_diffusion.py b/ldm/models/diffusion/shared_invokeai_diffusion.py index cddddd3e860..5c53ee79d6a 100644 --- a/ldm/models/diffusion/shared_invokeai_diffusion.py +++ b/ldm/models/diffusion/shared_invokeai_diffusion.py @@ -13,6 +13,7 @@ restore_default_cross_attention, override_cross_attention, Context, get_cross_attention_modules, \ CrossAttentionType, SwapCrossAttnContext from ldm.models.diffusion.cross_attention_map_saving import AttentionMapSaver +from ldm.modules.lora_manager import LoraManager ModelForwardCallback: TypeAlias = Union[ # x, t, conditioning, Optional[cross-attention kwargs] @@ -51,7 +52,7 @@ def wants_cross_attention_control(self): return self.cross_attention_control_args is not None - def __init__(self, model, model_forward_callback: ModelForwardCallback, + def __init__(self, model, model_forward_callback: ModelForwardCallback, lora_manager: LoraManager, is_running_diffusers: bool=False, ): """ @@ -64,6 +65,7 @@ def __init__(self, model, model_forward_callback: ModelForwardCallback, self.model_forward_callback = model_forward_callback self.cross_attention_control_context = None self.sequential_guidance = Globals.sequential_guidance + self.lora_manager = lora_manager @contextmanager def custom_attention_context(self, @@ -71,6 +73,8 @@ def custom_attention_context(self, step_count: int): do_swap = extra_conditioning_info is not None and extra_conditioning_info.wants_cross_attention_control old_attn_processor = None + if self.lora_manager: + self.lora_manager.load_loras() if do_swap: old_attn_processor = self.override_cross_attention(extra_conditioning_info, step_count=step_count) @@ -82,6 +86,7 @@ def custom_attention_context(self, # TODO resuscitate attention map saving #self.remove_attention_map_saving() + def override_cross_attention(self, conditioning: ExtraConditioningInfo, step_count: int) -> Dict[str, AttnProcessor]: """ setup cross attention .swap control. for diffusers this replaces the attention processor, so diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index 5c4d52d01f2..776de91d4d2 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -369,11 +369,15 @@ def __del__(self): class LoraManager: + models: list[str] + def __init__(self, pipe): self.lora_path = Path(global_models_dir(), 'lora') self.unet = pipe.unet + self.text_encoder = pipe.text_encoder # Legacy class handles lora not generated through diffusers self.legacy = LegacyLoraManager(pipe, self.lora_path) + self.models = [] def apply_lora_model(self, name): path = Path(self.lora_path, name) @@ -385,10 +389,17 @@ def apply_lora_model(self, name): else: print(f">> Unable to find valid LoRA at: {path}") - def load_lora_compel(self, lora_weights: list): + def set_lora_model(self, name): + self.models.append(name) + + def set_loras_compel(self, lora_weights: list): if len(lora_weights) > 0: for lora in lora_weights: - self.apply_lora_model(lora.model) + self.set_lora_model(lora.model) + + def load_loras(self): + for name in self.models: + self.apply_lora_model(name) # Legacy functions, to pipe to LoraLegacyManager def configure_prompt_legacy(self, prompt: str) -> str: From 68a3132d81807c9bd4c2ce8842365ce097c00785 Mon Sep 17 00:00:00 2001 From: Jordan Date: Thu, 23 Feb 2023 17:41:20 -0700 Subject: [PATCH 26/30] move legacy lora manager to its own file --- ldm/generate.py | 3 +- ldm/modules/legacy_lora_manager.py | 372 +++++++++++++++++++++++++++++ ldm/modules/lora_manager.py | 368 +--------------------------- 3 files changed, 376 insertions(+), 367 deletions(-) create mode 100644 ldm/modules/legacy_lora_manager.py diff --git a/ldm/generate.py b/ldm/generate.py index 9a7badb42e1..bcda1289410 100644 --- a/ldm/generate.py +++ b/ldm/generate.py @@ -488,7 +488,8 @@ def process_image(image,seed): self.sampler_name = sampler_name self._set_sampler() - # To try and load lora not trained through diffusers + # To try and load LoRA not trained through diffusers + # To be removed once support for diffusers LoRA weights is high enough if self.model.lora_manager: prompt = self.model.lora_manager.configure_prompt_legacy(prompt) # lora MUST process prompt before conditioning diff --git a/ldm/modules/legacy_lora_manager.py b/ldm/modules/legacy_lora_manager.py new file mode 100644 index 00000000000..a342f944934 --- /dev/null +++ b/ldm/modules/legacy_lora_manager.py @@ -0,0 +1,372 @@ +import re +from pathlib import Path +from ldm.invoke.devices import choose_torch_device +from safetensors.torch import load_file +import torch +from torch.utils.hooks import RemovableHandle +from diffusers.models import UNet2DConditionModel +from transformers import CLIPTextModel + +''' +This module supports loading LoRA weights trained with https://github.com/kohya-ss/sd-scripts +To be removed once support for diffusers LoRA weights is well supported +''' + + +class LoRALayer: + lora_name: str + name: str + scale: float + up: torch.nn.Module + down: torch.nn.Module + + def __init__(self, lora_name: str, name: str, rank=4, alpha=1.0): + self.lora_name = lora_name + self.name = name + self.scale = alpha / rank + + +class LoRAModuleWrapper: + unet: UNet2DConditionModel + text_encoder: CLIPTextModel + hooks: list[RemovableHandle] + + def __init__(self, unet, text_encoder): + self.unet = unet + self.text_encoder = text_encoder + self.hooks = [] + self.text_modules = None + self.unet_modules = None + + self.applied_loras = {} + self.loaded_loras = {} + + self.UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"] + self.TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"] + self.LORA_PREFIX_UNET = 'lora_unet' + self.LORA_PREFIX_TEXT_ENCODER = 'lora_te' + + self.re_digits = re.compile(r"\d+") + self.re_unet_transformer_attn_blocks = re.compile( + r"lora_unet_(.+)_blocks_(\d+)_attentions_(\d+)_transformer_blocks_(\d+)_attn(\d+)_(.+).(weight|alpha)" + ) + self.re_unet_mid_blocks = re.compile( + r"lora_unet_mid_block_attentions_(\d+)_(.+).(weight|alpha)" + ) + self.re_unet_transformer_blocks = re.compile( + r"lora_unet_(.+)_blocks_(\d+)_attentions_(\d+)_transformer_blocks_(\d+)_(.+).(weight|alpha)" + ) + self.re_unet_mid_transformer_blocks = re.compile( + r"lora_unet_mid_block_attentions_(\d+)_transformer_blocks_(\d+)_(.+).(weight|alpha)" + ) + self.re_unet_norm_blocks = re.compile( + r"lora_unet_(.+)_blocks_(\d+)_attentions_(\d+)_(.+).(weight|alpha)" + ) + self.re_out = re.compile(r"to_out_(\d+)") + self.re_processor_weight = re.compile(r"(.+)_(\d+)_(.+)") + self.re_processor_alpha = re.compile(r"(.+)_(\d+)") + + def find_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> dict[str, torch.nn.Module]: + mapping = {} + for name, module in root_module.named_modules(): + if module.__class__.__name__ in target_replace_modules: + for child_name, child_module in module.named_modules(): + layer_type = child_module.__class__.__name__ + if layer_type == "Linear" or (layer_type == "Conv2d" and child_module.kernel_size == (1, 1)): + lora_name = prefix + '.' + name + '.' + child_name + lora_name = lora_name.replace('.', '_') + mapping[lora_name] = child_module + self.apply_module_forward(child_module, lora_name) + return mapping + + if self.text_modules is None: + self.text_modules = find_modules( + self.LORA_PREFIX_TEXT_ENCODER, + text_encoder, + self.TEXT_ENCODER_TARGET_REPLACE_MODULE + ) + + if self.unet_modules is None: + self.unet_modules = find_modules( + self.LORA_PREFIX_UNET, + unet, + self.UNET_TARGET_REPLACE_MODULE + ) + + def convert_key_to_diffusers(self, key): + def match(match_list, regex, subject): + r = re.match(regex, subject) + if not r: + return False + + match_list.clear() + match_list.extend([int(x) if re.match(self.re_digits, x) else x for x in r.groups()]) + return True + + m = [] + + def get_front_block(first, second, third, fourth=None): + if first == "mid": + b_type = f"mid_block" + else: + b_type = f"{first}_blocks.{second}" + + if fourth is None: + return f"{b_type}.attentions.{third}" + + return f"{b_type}.attentions.{third}.transformer_blocks.{fourth}" + + def get_back_block(first, second, third): + second = second.replace(".lora_", "_lora.") + if third == "weight": + bm = [] + if match(bm, self.re_processor_weight, second): + s_bm = bm[2].split('.') + s_front = f"{bm[0]}_{s_bm[0]}" + s_back = f"{s_bm[1]}" + if int(bm[1]) == 0: + second = f"{s_front}.{s_back}" + else: + second = f"{s_front}.{bm[1]}.{s_back}" + elif third == "alpha": + bma = [] + if match(bma, self.re_processor_alpha, second): + if int(bma[1]) == 0: + second = f"{bma[0]}" + else: + second = f"{bma[0]}.{bma[1]}" + + if first is None: + return f"processor.{second}.{third}" + + return f"attn{first}.processor.{second}.{third}" + + if match(m, self.re_unet_transformer_attn_blocks, key): + return f"{get_front_block(m[0], m[1], m[2], m[3])}.{get_back_block(m[4], m[5], m[6])}" + + if match(m, self.re_unet_transformer_blocks, key): + return f"{get_front_block(m[0], m[1], m[2], m[3])}.{get_back_block(None, m[4], m[5])}" + + if match(m, self.re_unet_mid_transformer_blocks, key): + return f"{get_front_block('mid', None, m[0], m[1])}.{get_back_block(None, m[2], m[3])}" + + if match(m, self.re_unet_norm_blocks, key): + return f"{get_front_block(m[0], m[1], m[2])}.{get_back_block(None, m[3], m[4])}" + + if match(m, self.re_unet_mid_blocks, key): + return f"{get_front_block('mid', None, m[0])}.{get_back_block(None, m[1], m[2])}" + + return key + + def lora_forward_hook(self, name): + wrapper = self + + def lora_forward(module, input_h, output): + if len(wrapper.loaded_loras) == 0: + return output + + for lora in wrapper.applied_loras.values(): + layer = lora.layers.get(name, None) + if layer is None: + continue + output = output + layer.up(layer.down(*input_h)) * lora.multiplier * layer.scale + return output + + return lora_forward + + def apply_module_forward(self, module, name): + handle = module.register_forward_hook(self.lora_forward_hook(name)) + self.hooks.append(handle) + + def clear_hooks(self): + for hook in self.hooks: + hook.remove() + + self.hooks.clear() + + def clear_applied_loras(self): + self.applied_loras.clear() + + def clear_loaded_loras(self): + self.loaded_loras.clear() + + def __del__(self): + self.clear_hooks() + self.clear_applied_loras() + self.clear_loaded_loras() + del self.text_modules + del self.unet_modules + del self.hooks + + +class LoRA: + name: str + layers: dict[str, LoRALayer] + device: torch.device + dtype: torch.dtype + wrapper: LoRAModuleWrapper + multiplier: float + + def __init__(self, name: str, device, dtype, wrapper, multiplier=1.0): + self.name = name + self.layers = {} + self.multiplier = multiplier + self.device = device + self.dtype = dtype + self.wrapper = wrapper + self.rank = None + self.alpha = None + + def load_from_dict(self, state_dict): + for key, value in state_dict.items(): + stem, leaf = key.split(".", 1) + + if leaf.endswith("alpha"): + if self.alpha is None: + self.alpha = value.item() + continue + + if stem.startswith(self.wrapper.LORA_PREFIX_TEXT_ENCODER): + wrapped = self.wrapper.text_modules.get(stem, None) + if wrapped is None: + print(f">> Missing layer: {stem}") + continue + + if self.rank is None and leaf == 'lora_down.weight' and len(value.size()) == 2: + self.rank = value.shape[0] + self.load_lora_layer(stem, leaf, value, wrapped) + continue + elif stem.startswith(self.wrapper.LORA_PREFIX_UNET): + wrapped = self.wrapper.unet_modules.get(stem, None) + if wrapped is None: + print(f">> Missing layer: {stem}") + continue + + if self.rank is None and leaf == 'lora_down.weight' and len(value.size()) == 2: + self.rank = value.shape[0] + self.load_lora_layer(stem, leaf, value, wrapped) + continue + else: + continue + + def load_lora_layer(self, stem: str, leaf: str, value, wrapped: torch.nn.Module): + layer = self.layers.get(stem, None) + if layer is None: + layer = LoRALayer(self.name, stem, self.rank, self.alpha) + self.layers[stem] = layer + + if type(wrapped) == torch.nn.Linear: + module = torch.nn.Linear(value.shape[1], value.shape[0], bias=False) + elif type(wrapped) == torch.nn.Conv2d: + module = torch.nn.Conv2d(value.shape[1], value.shape[0], (1, 1), bias=False) + else: + print(f">> Encountered unknown lora layer module in {self.name}: {type(value).__name__}") + return + + with torch.no_grad(): + module.weight.copy_(value) + + module.to(device=self.device, dtype=self.dtype) + + if leaf == "lora_up.weight": + layer.up = module + elif leaf == "lora_down.weight": + layer.down = module + else: + print(f">> Encountered unknown layer in lora {self.name}: {leaf}") + return + + +class LegacyLoraManager: + def __init__(self, pipe, lora_path): + self.unet = pipe.unet + self.lora_path = lora_path + self.wrapper = LoRAModuleWrapper(pipe.unet, pipe.text_encoder) + self.text_encoder = pipe.text_encoder + self.device = torch.device(choose_torch_device()) + self.dtype = pipe.unet.dtype + self.loras_to_load = {} + + def load_lora_module(self, name, path_file, multiplier: float = 1.0): + # can be used instead to load through diffusers, once enough support is added + # lora = load_lora_attn(name, path_file, self.wrapper, multiplier) + + print(f">> Loading lora {name} from {path_file}") + if path_file.suffix == '.safetensors': + checkpoint = load_file(path_file.absolute().as_posix(), device='cpu') + else: + checkpoint = torch.load(path_file, map_location='cpu') + + lora = LoRA(name, self.device, self.dtype, self.wrapper, multiplier) + lora.load_from_dict(checkpoint) + self.wrapper.loaded_loras[name] = lora + + return lora + + def configure_prompt(self, prompt: str) -> str: + self.clear_loras() + + # lora_match = re.compile(r"]+)>") + lora_match = re.compile(r"withLoraLegacy\(([a-zA-Z\,\d]+)\)") + + for match in re.findall(lora_match, prompt): + # match = match.split(':') + match = match.split(',') + name = match[0].strip() + + mult = 1.0 + if len(match) == 2: + mult = float(match[1].strip()) + + self.set_lora(name, mult) + + # remove lora and return prompt to avoid the lora prompt causing issues in inference + return re.sub(lora_match, "", prompt) + + def apply_lora_model(self, name, mult: float = 1.0): + path_file = Path(self.lora_path, f'{name}.ckpt') + if Path(self.lora_path, f'{name}.safetensors').exists(): + path_file = Path(self.lora_path, f'{name}.safetensors') + + if not path_file.exists(): + print(f">> Unable to find lora: {name}") + return + + lora = self.wrapper.loaded_loras.get(name, None) + if lora is None: + lora = self.load_lora_module(name, path_file, mult) + + lora.multiplier = mult + self.wrapper.applied_loras[name] = lora + + def load_lora(self): + for name, multiplier in self.loras_to_load.items(): + self.apply_lora_model(name, multiplier) + + def unload_applied_loras(self, loras_to_load): + # unload any lora's not defined by loras_to_load + for name in list(self.wrapper.applied_loras.keys()): + if name not in loras_to_load: + self.unload_applied_lora(name) + + def unload_applied_lora(self, lora_name: str): + if lora_name in self.wrapper.applied_loras: + del self.wrapper.applied_loras[lora_name] + + def unload_lora(self, lora_name: str): + if lora_name in self.wrapper.loaded_loras: + del self.wrapper.loaded_loras[lora_name] + + def set_lora(self, name, multiplier: float = 1.0): + self.loras_to_load[name] = multiplier + + # update the multiplier if the lora was already loaded + if name in self.wrapper.loaded_loras: + self.wrapper.loaded_loras[name].multiplier = multiplier + + def clear_loras(self): + self.loras_to_load = {} + self.wrapper.clear_applied_loras() + + def __del__(self): + del self.loras_to_load diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index 776de91d4d2..0ea078c5257 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -1,371 +1,6 @@ -import re from pathlib import Path from ldm.invoke.globals import global_models_dir -from ldm.invoke.devices import choose_torch_device -from safetensors.torch import load_file -import torch -from torch.utils.hooks import RemovableHandle -from diffusers.models import UNet2DConditionModel -from transformers import CLIPTextModel - - -class LoRALayer: - lora_name: str - name: str - scale: float - up: torch.nn.Module - down: torch.nn.Module - - def __init__(self, lora_name: str, name: str, rank=4, alpha=1.0): - self.lora_name = lora_name - self.name = name - self.scale = alpha / rank - - -class LoRAModuleWrapper: - unet: UNet2DConditionModel - text_encoder: CLIPTextModel - hooks: list[RemovableHandle] - - def __init__(self, unet, text_encoder): - self.unet = unet - self.text_encoder = text_encoder - self.hooks = [] - self.text_modules = None - self.unet_modules = None - - self.applied_loras = {} - self.loaded_loras = {} - - self.UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"] - self.TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"] - self.LORA_PREFIX_UNET = 'lora_unet' - self.LORA_PREFIX_TEXT_ENCODER = 'lora_te' - - self.re_digits = re.compile(r"\d+") - self.re_unet_transformer_attn_blocks = re.compile( - r"lora_unet_(.+)_blocks_(\d+)_attentions_(\d+)_transformer_blocks_(\d+)_attn(\d+)_(.+).(weight|alpha)" - ) - self.re_unet_mid_blocks = re.compile( - r"lora_unet_mid_block_attentions_(\d+)_(.+).(weight|alpha)" - ) - self.re_unet_transformer_blocks = re.compile( - r"lora_unet_(.+)_blocks_(\d+)_attentions_(\d+)_transformer_blocks_(\d+)_(.+).(weight|alpha)" - ) - self.re_unet_mid_transformer_blocks = re.compile( - r"lora_unet_mid_block_attentions_(\d+)_transformer_blocks_(\d+)_(.+).(weight|alpha)" - ) - self.re_unet_norm_blocks = re.compile( - r"lora_unet_(.+)_blocks_(\d+)_attentions_(\d+)_(.+).(weight|alpha)" - ) - self.re_out = re.compile(r"to_out_(\d+)") - self.re_processor_weight = re.compile(r"(.+)_(\d+)_(.+)") - self.re_processor_alpha = re.compile(r"(.+)_(\d+)") - - def find_modules(prefix, root_module: torch.nn.Module, target_replace_modules) -> dict[str, torch.nn.Module]: - mapping = {} - for name, module in root_module.named_modules(): - if module.__class__.__name__ in target_replace_modules: - for child_name, child_module in module.named_modules(): - layer_type = child_module.__class__.__name__ - if layer_type == "Linear" or (layer_type == "Conv2d" and child_module.kernel_size == (1, 1)): - lora_name = prefix + '.' + name + '.' + child_name - lora_name = lora_name.replace('.', '_') - mapping[lora_name] = child_module - self.apply_module_forward(child_module, lora_name) - return mapping - - if self.text_modules is None: - self.text_modules = find_modules( - self.LORA_PREFIX_TEXT_ENCODER, - text_encoder, - self.TEXT_ENCODER_TARGET_REPLACE_MODULE - ) - - if self.unet_modules is None: - self.unet_modules = find_modules( - self.LORA_PREFIX_UNET, - unet, - self.UNET_TARGET_REPLACE_MODULE - ) - - def convert_key_to_diffusers(self, key): - def match(match_list, regex, subject): - r = re.match(regex, subject) - if not r: - return False - - match_list.clear() - match_list.extend([int(x) if re.match(self.re_digits, x) else x for x in r.groups()]) - return True - - m = [] - - def get_front_block(first, second, third, fourth=None): - if first == "mid": - b_type = f"mid_block" - else: - b_type = f"{first}_blocks.{second}" - - if fourth is None: - return f"{b_type}.attentions.{third}" - - return f"{b_type}.attentions.{third}.transformer_blocks.{fourth}" - - def get_back_block(first, second, third): - second = second.replace(".lora_", "_lora.") - if third == "weight": - bm = [] - if match(bm, self.re_processor_weight, second): - s_bm = bm[2].split('.') - s_front = f"{bm[0]}_{s_bm[0]}" - s_back = f"{s_bm[1]}" - if int(bm[1]) == 0: - second = f"{s_front}.{s_back}" - else: - second = f"{s_front}.{bm[1]}.{s_back}" - elif third == "alpha": - bma = [] - if match(bma, self.re_processor_alpha, second): - if int(bma[1]) == 0: - second = f"{bma[0]}" - else: - second = f"{bma[0]}.{bma[1]}" - - if first is None: - return f"processor.{second}.{third}" - - return f"attn{first}.processor.{second}.{third}" - - if match(m, self.re_unet_transformer_attn_blocks, key): - return f"{get_front_block(m[0], m[1], m[2], m[3])}.{get_back_block(m[4], m[5], m[6])}" - - if match(m, self.re_unet_transformer_blocks, key): - return f"{get_front_block(m[0], m[1], m[2], m[3])}.{get_back_block(None, m[4], m[5])}" - - if match(m, self.re_unet_mid_transformer_blocks, key): - return f"{get_front_block('mid', None, m[0], m[1])}.{get_back_block(None, m[2], m[3])}" - - if match(m, self.re_unet_norm_blocks, key): - return f"{get_front_block(m[0], m[1], m[2])}.{get_back_block(None, m[3], m[4])}" - - if match(m, self.re_unet_mid_blocks, key): - return f"{get_front_block('mid', None, m[0])}.{get_back_block(None, m[1], m[2])}" - - return key - - def lora_forward_hook(self, name): - wrapper = self - - def lora_forward(module, input_h, output): - if len(wrapper.loaded_loras) == 0: - return output - - for lora in wrapper.applied_loras.values(): - layer = lora.layers.get(name, None) - if layer is None: - continue - output = output + layer.up(layer.down(*input_h)) * lora.multiplier * layer.scale - return output - - return lora_forward - - def apply_module_forward(self, module, name): - handle = module.register_forward_hook(self.lora_forward_hook(name)) - self.hooks.append(handle) - - def clear_hooks(self): - for hook in self.hooks: - hook.remove() - - self.hooks.clear() - - def clear_applied_loras(self): - self.applied_loras.clear() - - def clear_loaded_loras(self): - self.loaded_loras.clear() - - def __del__(self): - self.clear_hooks() - self.clear_applied_loras() - self.clear_loaded_loras() - del self.text_modules - del self.unet_modules - del self.hooks - - -class LoRA: - name: str - layers: dict[str, LoRALayer] - device: torch.device - dtype: torch.dtype - wrapper: LoRAModuleWrapper - multiplier: float - - def __init__(self, name: str, device, dtype, wrapper, multiplier=1.0): - self.name = name - self.layers = {} - self.multiplier = multiplier - self.device = device - self.dtype = dtype - self.wrapper = wrapper - self.rank = None - self.alpha = None - - def load_from_dict(self, state_dict): - for key, value in state_dict.items(): - stem, leaf = key.split(".", 1) - - if leaf.endswith("alpha"): - if self.alpha is None: - self.alpha = value.item() - continue - - if stem.startswith(self.wrapper.LORA_PREFIX_TEXT_ENCODER): - wrapped = self.wrapper.text_modules.get(stem, None) - if wrapped is None: - print(f">> Missing layer: {stem}") - continue - - if self.rank is None and leaf == 'lora_down.weight' and len(value.size()) == 2: - self.rank = value.shape[0] - self.load_lora_layer(stem, leaf, value, wrapped) - continue - elif stem.startswith(self.wrapper.LORA_PREFIX_UNET): - wrapped = self.wrapper.unet_modules.get(stem, None) - if wrapped is None: - print(f">> Missing layer: {stem}") - continue - - if self.rank is None and leaf == 'lora_down.weight' and len(value.size()) == 2: - self.rank = value.shape[0] - self.load_lora_layer(stem, leaf, value, wrapped) - continue - else: - continue - - def load_lora_layer(self, stem: str, leaf: str, value, wrapped: torch.nn.Module): - layer = self.layers.get(stem, None) - if layer is None: - layer = LoRALayer(self.name, stem, self.rank, self.alpha) - self.layers[stem] = layer - - if type(wrapped) == torch.nn.Linear: - module = torch.nn.Linear(value.shape[1], value.shape[0], bias=False) - elif type(wrapped) == torch.nn.Conv2d: - module = torch.nn.Conv2d(value.shape[1], value.shape[0], (1, 1), bias=False) - else: - print(f">> Encountered unknown lora layer module in {self.name}: {type(value).__name__}") - return - - with torch.no_grad(): - module.weight.copy_(value) - - module.to(device=self.device, dtype=self.dtype) - - if leaf == "lora_up.weight": - layer.up = module - elif leaf == "lora_down.weight": - layer.down = module - else: - print(f">> Encountered unknown layer in lora {self.name}: {leaf}") - return - - -class LegacyLoraManager: - def __init__(self, pipe, lora_path): - self.unet = pipe.unet - self.lora_path = lora_path - self.wrapper = LoRAModuleWrapper(pipe.unet, pipe.text_encoder) - self.text_encoder = pipe.text_encoder - self.device = torch.device(choose_torch_device()) - self.dtype = pipe.unet.dtype - self.loras_to_load = {} - - def load_lora_module(self, name, path_file, multiplier: float = 1.0): - # can be used instead to load through diffusers, once enough support is added - # lora = load_lora_attn(name, path_file, self.wrapper, multiplier) - - print(f">> Loading lora {name} from {path_file}") - if path_file.suffix == '.safetensors': - checkpoint = load_file(path_file.absolute().as_posix(), device='cpu') - else: - checkpoint = torch.load(path_file, map_location='cpu') - - lora = LoRA(name, self.device, self.dtype, self.wrapper, multiplier) - lora.load_from_dict(checkpoint) - self.wrapper.loaded_loras[name] = lora - - return lora - - def configure_prompt(self, prompt: str) -> str: - self.clear_loras() - - # lora_match = re.compile(r"]+)>") - lora_match = re.compile(r"withLoraLegacy\(([a-zA-Z\,\d]+)\)") - - for match in re.findall(lora_match, prompt): - # match = match.split(':') - match = match.split(',') - name = match[0].strip() - - mult = 1.0 - if len(match) == 2: - mult = float(match[1].strip()) - - self.set_lora(name, mult) - - # remove lora and return prompt to avoid the lora prompt causing issues in inference - return re.sub(lora_match, "", prompt) - - def apply_lora_model(self, name, mult: float = 1.0): - path_file = Path(self.lora_path, f'{name}.ckpt') - if Path(self.lora_path, f'{name}.safetensors').exists(): - path_file = Path(self.lora_path, f'{name}.safetensors') - - if not path_file.exists(): - print(f">> Unable to find lora: {name}") - return - - lora = self.wrapper.loaded_loras.get(name, None) - if lora is None: - lora = self.load_lora_module(name, path_file, mult) - - lora.multiplier = mult - self.wrapper.applied_loras[name] = lora - - def load_lora(self): - for name, multiplier in self.loras_to_load.items(): - self.apply_lora_model(name, multiplier) - - def unload_applied_loras(self, loras_to_load): - # unload any lora's not defined by loras_to_load - for name in list(self.wrapper.applied_loras.keys()): - if name not in loras_to_load: - self.unload_applied_lora(name) - - def unload_applied_lora(self, lora_name: str): - if lora_name in self.wrapper.applied_loras: - del self.wrapper.applied_loras[lora_name] - - def unload_lora(self, lora_name: str): - if lora_name in self.wrapper.loaded_loras: - del self.wrapper.loaded_loras[lora_name] - - def set_lora(self, name, multiplier: float = 1.0): - self.loras_to_load[name] = multiplier - - # update the multiplier if the lora was already loaded - if name in self.wrapper.loaded_loras: - self.wrapper.loaded_loras[name].multiplier = multiplier - - def clear_loras(self): - self.loras_to_load = {} - self.wrapper.clear_applied_loras() - - def __del__(self): - del self.loras_to_load +from .legacy_lora_manager import LegacyLoraManager class LoraManager: @@ -402,6 +37,7 @@ def load_loras(self): self.apply_lora_model(name) # Legacy functions, to pipe to LoraLegacyManager + # To be removed once support for diffusers LoRA weights is high enough def configure_prompt_legacy(self, prompt: str) -> str: return self.legacy.configure_prompt(prompt) From 4ce8b1ba21b94608f40dd4f582dc12729b3994f7 Mon Sep 17 00:00:00 2001 From: Jordan Date: Thu, 23 Feb 2023 19:27:45 -0700 Subject: [PATCH 27/30] setup cross conditioning for lora --- ldm/invoke/conditioning.py | 6 ++- ldm/invoke/generator/diffusers_pipeline.py | 5 +- ldm/invoke/globals.py | 4 ++ .../diffusion/shared_invokeai_diffusion.py | 32 ++++++++--- ldm/modules/lora_manager.py | 54 +++++++++++-------- 5 files changed, 68 insertions(+), 33 deletions(-) diff --git a/ldm/invoke/conditioning.py b/ldm/invoke/conditioning.py index 32b19fab655..d4fc978e904 100644 --- a/ldm/invoke/conditioning.py +++ b/ldm/invoke/conditioning.py @@ -55,12 +55,13 @@ def get_uc_and_c_and_ec(prompt_string, model, log_tokens=False, skip_normalize_l positive_prompt_string, negative_prompt_string = split_prompt_to_positive_and_negative(prompt_string) legacy_blend = try_parse_legacy_blend(positive_prompt_string, skip_normalize_legacy_blend) positive_prompt: FlattenedPrompt|Blend + lora_conditions = None if legacy_blend is not None: positive_prompt = legacy_blend else: positive_prompt = Compel.parse_prompt_string(positive_prompt_string) if model.lora_manager: - model.lora_manager.set_loras_compel(positive_prompt.lora_weights) + lora_conditions = model.lora_manager.set_loras_conditions(positive_prompt.lora_weights) negative_prompt: FlattenedPrompt|Blend = Compel.parse_prompt_string(negative_prompt_string) if log_tokens or getattr(Globals, "log_tokenization", False): @@ -73,7 +74,8 @@ def get_uc_and_c_and_ec(prompt_string, model, log_tokens=False, skip_normalize_l ec = InvokeAIDiffuserComponent.ExtraConditioningInfo(tokens_count_including_eos_bos=tokens_count, cross_attention_control_args=options.get( - 'cross_attention_control', None)) + 'cross_attention_control', None), + lora_conditions=lora_conditions) return uc, c, ec diff --git a/ldm/invoke/generator/diffusers_pipeline.py b/ldm/invoke/generator/diffusers_pipeline.py index 7fd4b5880ed..46034b845b0 100644 --- a/ldm/invoke/generator/diffusers_pipeline.py +++ b/ldm/invoke/generator/diffusers_pipeline.py @@ -291,10 +291,7 @@ def __init__( feature_extractor=feature_extractor, ) self.lora_manager = LoraManager(self) - self.invokeai_diffuser = InvokeAIDiffuserComponent(self.unet, - self._unet_forward, - self.lora_manager, - is_running_diffusers=True) + self.invokeai_diffuser = InvokeAIDiffuserComponent(self.unet, self._unet_forward, is_running_diffusers=True) use_full_precision = (precision == 'float32' or precision == 'autocast') self.textual_inversion_manager = TextualInversionManager(tokenizer=self.tokenizer, text_encoder=self.text_encoder, diff --git a/ldm/invoke/globals.py b/ldm/invoke/globals.py index e47b5c059e5..2f4d9fdd779 100644 --- a/ldm/invoke/globals.py +++ b/ldm/invoke/globals.py @@ -23,6 +23,7 @@ Globals.initfile = 'invokeai.init' Globals.models_file = 'models.yaml' Globals.models_dir = 'models' +Globals.lora_models_dir = 'lora' Globals.config_dir = 'configs' Globals.autoscan_dir = 'weights' Globals.converted_ckpts_dir = 'converted_ckpts' @@ -75,6 +76,9 @@ def global_config_dir()->Path: def global_models_dir()->Path: return Path(Globals.root, Globals.models_dir) +def global_lora_models_dir()->Path: + return Path(global_models_dir(), Globals.lora_models_dir) + def global_autoscan_dir()->Path: return Path(Globals.root, Globals.autoscan_dir) diff --git a/ldm/models/diffusion/shared_invokeai_diffusion.py b/ldm/models/diffusion/shared_invokeai_diffusion.py index 5c53ee79d6a..8ee87c53c1f 100644 --- a/ldm/models/diffusion/shared_invokeai_diffusion.py +++ b/ldm/models/diffusion/shared_invokeai_diffusion.py @@ -13,7 +13,7 @@ restore_default_cross_attention, override_cross_attention, Context, get_cross_attention_modules, \ CrossAttentionType, SwapCrossAttnContext from ldm.models.diffusion.cross_attention_map_saving import AttentionMapSaver -from ldm.modules.lora_manager import LoraManager +from ldm.modules.lora_manager import LoraCondition ModelForwardCallback: TypeAlias = Union[ # x, t, conditioning, Optional[cross-attention kwargs] @@ -46,13 +46,21 @@ class ExtraConditioningInfo: tokens_count_including_eos_bos: int cross_attention_control_args: Optional[Arguments] = None + lora_conditions: Optional[list[LoraCondition]] = None @property def wants_cross_attention_control(self): return self.cross_attention_control_args is not None + @property + def has_lora_conditions(self): + return self.lora_conditions is not None + + @property + def should_do_swap(self): + return self.wants_cross_attention_control or self.has_lora_conditions - def __init__(self, model, model_forward_callback: ModelForwardCallback, lora_manager: LoraManager, + def __init__(self, model, model_forward_callback: ModelForwardCallback, is_running_diffusers: bool=False, ): """ @@ -65,16 +73,13 @@ def __init__(self, model, model_forward_callback: ModelForwardCallback, lora_man self.model_forward_callback = model_forward_callback self.cross_attention_control_context = None self.sequential_guidance = Globals.sequential_guidance - self.lora_manager = lora_manager @contextmanager def custom_attention_context(self, extra_conditioning_info: Optional[ExtraConditioningInfo], step_count: int): - do_swap = extra_conditioning_info is not None and extra_conditioning_info.wants_cross_attention_control + do_swap = extra_conditioning_info is not None and extra_conditioning_info.should_do_swap old_attn_processor = None - if self.lora_manager: - self.lora_manager.load_loras() if do_swap: old_attn_processor = self.override_cross_attention(extra_conditioning_info, step_count=step_count) @@ -93,6 +98,21 @@ def override_cross_attention(self, conditioning: ExtraConditioningInfo, step_cou the previous attention processor is returned so that the caller can restore it later. """ self.conditioning = conditioning + + # If other modules do not want cross_attention_control then we should bypass setting up Context + old_attn_processors = None + if not self.conditioning.wants_cross_attention_control: + old_attn_processors = self.model.attn_processors + + # Load lora conditions into the model + if self.conditioning.has_lora_conditions: + for condition in self.conditioning.lora_conditions: + condition(self.model) + + # return old_attn_processors if there is nothing further to do here + if not self.conditioning.wants_cross_attention_control: + return old_attn_processors + self.cross_attention_control_context = Context( arguments=self.conditioning.cross_attention_control_args, step_count=step_count diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index 0ea078c5257..89ae21c8133 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -1,40 +1,52 @@ from pathlib import Path -from ldm.invoke.globals import global_models_dir +from ldm.invoke.globals import global_lora_models_dir from .legacy_lora_manager import LegacyLoraManager -class LoraManager: - models: list[str] +class LoraCondition: + name: str + weight: float - def __init__(self, pipe): - self.lora_path = Path(global_models_dir(), 'lora') - self.unet = pipe.unet - self.text_encoder = pipe.text_encoder - # Legacy class handles lora not generated through diffusers - self.legacy = LegacyLoraManager(pipe, self.lora_path) - self.models = [] + def __init__(self, name, weight: float = 1.0): + self.name = name + self.weight = weight - def apply_lora_model(self, name): - path = Path(self.lora_path, name) + def __call__(self, model): + path = Path(global_lora_models_dir(), self.name) file = Path(path, "pytorch_lora_weights.bin") if path.is_dir() and file.is_file(): - print(f">> Loading LoRA: {path}") - self.unet.load_attn_procs(path.absolute().as_posix()) + if model.load_attn_procs: + print(f">> Loading LoRA: {path}") + model.load_attn_procs(path.absolute().as_posix()) + else: + print(f">> Invalid Model to load LoRA") else: print(f">> Unable to find valid LoRA at: {path}") - def set_lora_model(self, name): - self.models.append(name) - def set_loras_compel(self, lora_weights: list): +class LoraManager: + conditions: list[LoraCondition] + + def __init__(self, pipe): + self.unet = pipe.unet + self.text_encoder = pipe.text_encoder + # Legacy class handles lora not generated through diffusers + self.legacy = LegacyLoraManager(pipe, global_lora_models_dir()) + self.conditions = [] + + def set_lora_model(self, name, weight: float = 1.0): + self.conditions.append(LoraCondition(name, weight)) + + def set_loras_conditions(self, lora_weights: list): if len(lora_weights) > 0: for lora in lora_weights: - self.set_lora_model(lora.model) + self.set_lora_model(lora.model, lora.weight) + + if len(self.conditions) > 0: + return self.conditions - def load_loras(self): - for name in self.models: - self.apply_lora_model(name) + return None # Legacy functions, to pipe to LoraLegacyManager # To be removed once support for diffusers LoRA weights is high enough From 523e44ccfe464a13056d5aade7fefe845e1a7fb0 Mon Sep 17 00:00:00 2001 From: Jordan Date: Fri, 24 Feb 2023 01:32:09 -0700 Subject: [PATCH 28/30] simplify manager --- ldm/modules/lora_manager.py | 30 ++++++++++++++---------------- 1 file changed, 14 insertions(+), 16 deletions(-) diff --git a/ldm/modules/lora_manager.py b/ldm/modules/lora_manager.py index 89ae21c8133..885d82b53bd 100644 --- a/ldm/modules/lora_manager.py +++ b/ldm/modules/lora_manager.py @@ -13,12 +13,16 @@ def __init__(self, name, weight: float = 1.0): def __call__(self, model): path = Path(global_lora_models_dir(), self.name) - file = Path(path, "pytorch_lora_weights.bin") - if path.is_dir() and file.is_file(): + # TODO: make model able to load from huggingface, rather then just local files + if path.is_dir(): if model.load_attn_procs: - print(f">> Loading LoRA: {path}") - model.load_attn_procs(path.absolute().as_posix()) + file = Path(path, "pytorch_lora_weights.bin") + if file.is_file(): + print(f">> Loading LoRA: {path}") + model.load_attn_procs(path.absolute().as_posix()) + else: + print(f">> Unable to find valid LoRA at: {path}") else: print(f">> Invalid Model to load LoRA") else: @@ -26,25 +30,19 @@ def __call__(self, model): class LoraManager: - conditions: list[LoraCondition] - def __init__(self, pipe): - self.unet = pipe.unet - self.text_encoder = pipe.text_encoder # Legacy class handles lora not generated through diffusers self.legacy = LegacyLoraManager(pipe, global_lora_models_dir()) - self.conditions = [] - - def set_lora_model(self, name, weight: float = 1.0): - self.conditions.append(LoraCondition(name, weight)) - def set_loras_conditions(self, lora_weights: list): + @staticmethod + def set_loras_conditions(lora_weights: list): + conditions = [] if len(lora_weights) > 0: for lora in lora_weights: - self.set_lora_model(lora.model, lora.weight) + conditions.append(LoraCondition(lora.model, lora.weight)) - if len(self.conditions) > 0: - return self.conditions + if len(conditions) > 0: + return conditions return None From 7dbe027b180a9a32ae5d13666139dede94688fc5 Mon Sep 17 00:00:00 2001 From: Damian Stewart Date: Fri, 24 Feb 2023 12:46:57 +0100 Subject: [PATCH 29/30] tweaks and small refactors --- .../diffusion/cross_attention_control.py | 6 +-- ldm/models/diffusion/ddim.py | 2 +- ldm/models/diffusion/ksampler.py | 2 +- ldm/models/diffusion/plms.py | 2 +- .../diffusion/shared_invokeai_diffusion.py | 50 ++++++++----------- 5 files changed, 25 insertions(+), 37 deletions(-) diff --git a/ldm/models/diffusion/cross_attention_control.py b/ldm/models/diffusion/cross_attention_control.py index a34f22e683f..7a66448b1cb 100644 --- a/ldm/models/diffusion/cross_attention_control.py +++ b/ldm/models/diffusion/cross_attention_control.py @@ -289,10 +289,10 @@ def get_invokeai_attention_mem_efficient(self, q, k, v): -def restore_default_cross_attention(model, is_running_diffusers: bool, restore_attention_processor: Optional[AttnProcessor]=None): +def restore_default_cross_attention(model, is_running_diffusers: bool, processors_to_restore: Optional[AttnProcessor]=None): if is_running_diffusers: unet = model - unet.set_attn_processor(restore_attention_processor or CrossAttnProcessor()) + unet.set_attn_processor(processors_to_restore or CrossAttnProcessor()) else: remove_attention_function(model) @@ -334,11 +334,9 @@ def override_cross_attention(model, context: Context, is_running_diffusers = Fal default_slice_size = 4 slice_size = next((p.slice_size for p in old_attn_processors.values() if type(p) is SlicedAttnProcessor), default_slice_size) unet.set_attn_processor(SlicedSwapCrossAttnProcesser(slice_size=slice_size)) - return old_attn_processors else: context.register_cross_attention_modules(model) inject_attention_function(model, context) - return None diff --git a/ldm/models/diffusion/ddim.py b/ldm/models/diffusion/ddim.py index 304009c1d3c..ebf3ea22b6e 100644 --- a/ldm/models/diffusion/ddim.py +++ b/ldm/models/diffusion/ddim.py @@ -19,7 +19,7 @@ def prepare_to_sample(self, t_enc, **kwargs): all_timesteps_count = kwargs.get('all_timesteps_count', t_enc) if extra_conditioning_info is not None and extra_conditioning_info.wants_cross_attention_control: - self.invokeai_diffuser.override_cross_attention(extra_conditioning_info, step_count = all_timesteps_count) + self.invokeai_diffuser.override_attention_processors(extra_conditioning_info, step_count = all_timesteps_count) else: self.invokeai_diffuser.restore_default_cross_attention() diff --git a/ldm/models/diffusion/ksampler.py b/ldm/models/diffusion/ksampler.py index f98ca8de21a..d5bc2b5d146 100644 --- a/ldm/models/diffusion/ksampler.py +++ b/ldm/models/diffusion/ksampler.py @@ -43,7 +43,7 @@ def prepare_to_sample(self, t_enc, **kwargs): extra_conditioning_info = kwargs.get('extra_conditioning_info', None) if extra_conditioning_info is not None and extra_conditioning_info.wants_cross_attention_control: - self.invokeai_diffuser.override_cross_attention(extra_conditioning_info, step_count = t_enc) + self.invokeai_diffuser.override_attention_processors(extra_conditioning_info, step_count = t_enc) else: self.invokeai_diffuser.restore_default_cross_attention() diff --git a/ldm/models/diffusion/plms.py b/ldm/models/diffusion/plms.py index 9edd3337803..aa8ccc1b5e7 100644 --- a/ldm/models/diffusion/plms.py +++ b/ldm/models/diffusion/plms.py @@ -21,7 +21,7 @@ def prepare_to_sample(self, t_enc, **kwargs): all_timesteps_count = kwargs.get('all_timesteps_count', t_enc) if extra_conditioning_info is not None and extra_conditioning_info.wants_cross_attention_control: - self.invokeai_diffuser.override_cross_attention(extra_conditioning_info, step_count = all_timesteps_count) + self.invokeai_diffuser.override_attention_processors(extra_conditioning_info, step_count = all_timesteps_count) else: self.invokeai_diffuser.restore_default_cross_attention() diff --git a/ldm/models/diffusion/shared_invokeai_diffusion.py b/ldm/models/diffusion/shared_invokeai_diffusion.py index 8ee87c53c1f..bee82329f78 100644 --- a/ldm/models/diffusion/shared_invokeai_diffusion.py +++ b/ldm/models/diffusion/shared_invokeai_diffusion.py @@ -56,9 +56,6 @@ def wants_cross_attention_control(self): def has_lora_conditions(self): return self.lora_conditions is not None - @property - def should_do_swap(self): - return self.wants_cross_attention_control or self.has_lora_conditions def __init__(self, model, model_forward_callback: ModelForwardCallback, is_running_diffusers: bool=False, @@ -78,11 +75,11 @@ def __init__(self, model, model_forward_callback: ModelForwardCallback, def custom_attention_context(self, extra_conditioning_info: Optional[ExtraConditioningInfo], step_count: int): - do_swap = extra_conditioning_info is not None and extra_conditioning_info.should_do_swap old_attn_processor = None - if do_swap: - old_attn_processor = self.override_cross_attention(extra_conditioning_info, - step_count=step_count) + if extra_conditioning_info.wants_cross_attention_control | extra_conditioning_info.has_lora_conditions: + old_attn_processor = self.override_attention_processors(extra_conditioning_info, + step_count=step_count) + try: yield None finally: @@ -92,41 +89,34 @@ def custom_attention_context(self, #self.remove_attention_map_saving() - def override_cross_attention(self, conditioning: ExtraConditioningInfo, step_count: int) -> Dict[str, AttnProcessor]: + def override_attention_processors(self, conditioning: ExtraConditioningInfo, step_count: int) -> Dict[str, AttnProcessor]: """ setup cross attention .swap control. for diffusers this replaces the attention processor, so the previous attention processor is returned so that the caller can restore it later. """ - self.conditioning = conditioning - - # If other modules do not want cross_attention_control then we should bypass setting up Context - old_attn_processors = None - if not self.conditioning.wants_cross_attention_control: - old_attn_processors = self.model.attn_processors + old_attn_processors = self.model.attn_processors # Load lora conditions into the model - if self.conditioning.has_lora_conditions: - for condition in self.conditioning.lora_conditions: + if conditioning.has_lora_conditions: + for condition in conditioning.lora_conditions: condition(self.model) - # return old_attn_processors if there is nothing further to do here - if not self.conditioning.wants_cross_attention_control: - return old_attn_processors + if conditioning.wants_cross_attention_control: + self.cross_attention_control_context = Context( + arguments=conditioning.cross_attention_control_args, + step_count=step_count + ) + override_cross_attention(self.model, + self.cross_attention_control_context, + is_running_diffusers=self.is_running_diffusers) + return old_attn_processors - self.cross_attention_control_context = Context( - arguments=self.conditioning.cross_attention_control_args, - step_count=step_count - ) - return override_cross_attention(self.model, - self.cross_attention_control_context, - is_running_diffusers=self.is_running_diffusers) - def restore_default_cross_attention(self, restore_attention_processor: Optional['AttnProcessor']=None): - self.conditioning = None + def restore_default_cross_attention(self, processors_to_restore: Optional[dict[str, 'AttnProcessor']]=None): self.cross_attention_control_context = None restore_default_cross_attention(self.model, is_running_diffusers=self.is_running_diffusers, - restore_attention_processor=restore_attention_processor) + processors_to_restore=processors_to_restore) def setup_attention_map_saving(self, saver: AttentionMapSaver): def callback(slice, dim, offset, slice_size, key): @@ -328,7 +318,7 @@ def _apply_cross_attention_controlled_conditioning__compvis(self, x:torch.Tensor #print("applying saved attention maps for", cross_attention_control_types_to_do) for ca_type in cross_attention_control_types_to_do: context.request_apply_saved_attention_maps(ca_type) - edited_conditioning = self.conditioning.cross_attention_control_args.edited_conditioning + edited_conditioning = context.arguments.edited_conditioning conditioned_next_x = self.model_forward_callback(x, sigma, edited_conditioning) context.clear_requests(cleanup=True) From d9c46277ea53652f7ca279c6064ea74450015e11 Mon Sep 17 00:00:00 2001 From: Jordan Date: Sat, 25 Feb 2023 20:21:20 -0700 Subject: [PATCH 30/30] add peft setup (need to install huggingface/peft) --- ldm/generate.py | 3 + ldm/invoke/conditioning.py | 7 +- ldm/invoke/generator/diffusers_pipeline.py | 2 + ldm/modules/peft_manager.py | 96 ++++++++++++++++++++++ 4 files changed, 107 insertions(+), 1 deletion(-) create mode 100644 ldm/modules/peft_manager.py diff --git a/ldm/generate.py b/ldm/generate.py index 1cb11d3606c..43d747afaed 100644 --- a/ldm/generate.py +++ b/ldm/generate.py @@ -533,6 +533,9 @@ def process_image(image,seed): log_tokens=self.log_tokenization, ) + if self.model.peft_manager: + self.model = self.model.peft_manager.load(self.model, self.model.unet.dtype) + init_image, mask_image = self._make_images( init_img, init_mask, diff --git a/ldm/invoke/conditioning.py b/ldm/invoke/conditioning.py index 2a0ab1ce024..02f8fc573ab 100644 --- a/ldm/invoke/conditioning.py +++ b/ldm/invoke/conditioning.py @@ -60,7 +60,12 @@ def get_uc_and_c_and_ec(prompt_string, model, log_tokens=False, skip_normalize_l positive_prompt = legacy_blend else: positive_prompt = Compel.parse_prompt_string(positive_prompt_string) - if model.lora_manager: + should_use_lora_manager = True + if model.peft_manager: + should_use_lora_manager = model.peft_manager.should_use(positive_prompt.lora_weights) + if not should_use_lora_manager: + model.peft_manager.set_loras(positive_prompt.lora_weights) + if model.lora_manager and should_use_lora_manager: lora_conditions = model.lora_manager.set_loras_conditions(positive_prompt.lora_weights) negative_prompt: FlattenedPrompt|Blend = Compel.parse_prompt_string(negative_prompt_string) diff --git a/ldm/invoke/generator/diffusers_pipeline.py b/ldm/invoke/generator/diffusers_pipeline.py index 46034b845b0..2dd90c8d1e1 100644 --- a/ldm/invoke/generator/diffusers_pipeline.py +++ b/ldm/invoke/generator/diffusers_pipeline.py @@ -30,6 +30,7 @@ from ldm.models.diffusion.shared_invokeai_diffusion import InvokeAIDiffuserComponent, PostprocessingSettings from ldm.modules.textual_inversion_manager import TextualInversionManager from ldm.modules.lora_manager import LoraManager +from ldm.modules.peft_manager import PeftManager from ..devices import normalize_device, CPU_DEVICE from ..offloading import LazilyLoadedModelGroup, FullyLoadedModelGroup, ModelGroup from ...models.diffusion.cross_attention_map_saving import AttentionMapSaver @@ -291,6 +292,7 @@ def __init__( feature_extractor=feature_extractor, ) self.lora_manager = LoraManager(self) + self.peft_manager = PeftManager() self.invokeai_diffuser = InvokeAIDiffuserComponent(self.unet, self._unet_forward, is_running_diffusers=True) use_full_precision = (precision == 'float32' or precision == 'autocast') self.textual_inversion_manager = TextualInversionManager(tokenizer=self.tokenizer, diff --git a/ldm/modules/peft_manager.py b/ldm/modules/peft_manager.py new file mode 100644 index 00000000000..939a4a78992 --- /dev/null +++ b/ldm/modules/peft_manager.py @@ -0,0 +1,96 @@ +from peft import LoraModel, LoraConfig, set_peft_model_state_dict +import torch +import json +from pathlib import Path +from ldm.invoke.globals import global_lora_models_dir + + +class LoraPeftModule: + def __init__(self, lora_dir, multiplier: float = 1.0): + self.lora_dir = lora_dir + self.multiplier = multiplier + self.config = self.load_config() + self.checkpoint = self.load_checkpoint() + + def load_config(self): + lora_config_file = Path(self.lora_dir, f'lora_config.json') + with open(lora_config_file, "r") as f: + return json.load(f) + + def load_checkpoint(self): + return torch.load(Path(self.lora_dir, f'lora.pt')) + + def unet(self, text_encoder): + lora_ds = { + k.replace("text_encoder_", ""): v for k, v in self.checkpoint.items() if "text_encoder_" in k + } + config = LoraConfig(**self.config["peft_config"]) + model = LoraModel(config, text_encoder) + set_peft_model_state_dict(model, lora_ds) + return model + + def text_encoder(self, unet): + lora_ds = { + k: v for k, v in self.checkpoint.items() if "text_encoder_" not in k + } + config = LoraConfig(**self.config["text_encoder_peft_config"]) + model = LoraModel(config, unet) + set_peft_model_state_dict(model, lora_ds) + return model + + def apply(self, pipe, dtype): + pipe.unet = self.unet(pipe.unet) + if "text_encoder_peft_config" in self.config: + pipe.text_encoder = self.text_encoder(pipe.text_encoder) + + if dtype in (torch.float16, torch.bfloat16): + pipe.unet.half() + pipe.text_encoder.half() + + return pipe + + +class PeftManager: + modules: list[LoraPeftModule] + + def __init__(self): + self.lora_path = global_lora_models_dir() + self.modules = [] + + def set_loras(self, lora_weights: list): + if len(lora_weights) > 0: + for lora in lora_weights: + self.add(lora.model, lora.weight) + + def add(self, name, multiplier: float = 1.0): + lora_dir = Path(self.lora_path, name) + + if lora_dir.exists(): + lora_config_file = Path(lora_dir, f'lora_config.json') + lora_checkpoint = Path(lora_dir, f'lora.pt') + + if lora_config_file.exists() and lora_checkpoint.exists(): + self.modules.append(LoraPeftModule(lora_dir, multiplier)) + return + + print(f">> Failed to load lora {name}") + + def load(self, pipe, dtype): + if len(self.modules) > 0: + for module in self.modules: + pipe = module.apply(pipe, dtype) + + return pipe + + # Simple check to allow previous functionality + def should_use(self, lora_weights: list): + if len(lora_weights) > 0: + for lora in lora_weights: + lora_dir = Path(self.lora_path, lora.model) + if lora_dir.exists(): + lora_config_file = Path(lora_dir, f'lora_config.json') + lora_checkpoint = Path(lora_dir, f'lora.pt') + if lora_config_file.exists() and lora_checkpoint.exists(): + return False + + return True