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 256f214b253..a038be9eee6 100644 --- a/ldm/generate.py +++ b/ldm/generate.py @@ -488,6 +488,13 @@ def process_image(image,seed): self.sampler_name = sampler_name self._set_sampler() + # 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 + self.model.lora_manager.load_lora_legacy() + # apply the concepts library to the prompt prompt = self.huggingface_concepts_library.replace_concepts_with_triggers( prompt, @@ -526,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, @@ -988,6 +998,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/conditioning.py b/ldm/invoke/conditioning.py index 7c654caf693..02f8fc573ab 100644 --- a/ldm/invoke/conditioning.py +++ b/ldm/invoke/conditioning.py @@ -55,10 +55,18 @@ 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) + 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) if log_tokens or getattr(Globals, "log_tokenization", False): @@ -71,7 +79,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 5e65cb5d138..2dd90c8d1e1 100644 --- a/ldm/invoke/generator/diffusers_pipeline.py +++ b/ldm/invoke/generator/diffusers_pipeline.py @@ -29,6 +29,8 @@ 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 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 @@ -289,11 +291,14 @@ def __init__( safety_checker=safety_checker, 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, text_encoder=self.text_encoder, full_precision=use_full_precision) + # InvokeAI's interface for text embeddings and whatnot self.embeddings_provider = EmbeddingsProvider( tokenizer=self.tokenizer, 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/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 cddddd3e860..bee82329f78 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 LoraCondition ModelForwardCallback: TypeAlias = Union[ # x, t, conditioning, Optional[cross-attention kwargs] @@ -45,11 +46,16 @@ 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 + def __init__(self, model, model_forward_callback: ModelForwardCallback, is_running_diffusers: bool=False, @@ -69,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.wants_cross_attention_control 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: @@ -82,26 +88,35 @@ 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]: + + 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 - 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 + old_attn_processors = self.model.attn_processors + + # Load lora conditions into the model + if conditioning.has_lora_conditions: + for condition in conditioning.lora_conditions: + condition(self.model) + + 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 + + + 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): @@ -303,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) 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 new file mode 100644 index 00000000000..885d82b53bd --- /dev/null +++ b/ldm/modules/lora_manager.py @@ -0,0 +1,55 @@ +from pathlib import Path +from ldm.invoke.globals import global_lora_models_dir +from .legacy_lora_manager import LegacyLoraManager + + +class LoraCondition: + name: str + weight: float + + def __init__(self, name, weight: float = 1.0): + self.name = name + self.weight = weight + + def __call__(self, model): + path = Path(global_lora_models_dir(), self.name) + + # TODO: make model able to load from huggingface, rather then just local files + if path.is_dir(): + if model.load_attn_procs: + 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: + print(f">> Unable to find valid LoRA at: {path}") + + +class LoraManager: + def __init__(self, pipe): + # Legacy class handles lora not generated through diffusers + self.legacy = LegacyLoraManager(pipe, global_lora_models_dir()) + + @staticmethod + def set_loras_conditions(lora_weights: list): + conditions = [] + if len(lora_weights) > 0: + for lora in lora_weights: + conditions.append(LoraCondition(lora.model, lora.weight)) + + if len(conditions) > 0: + return conditions + + return None + + # 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) + + def load_lora_legacy(self): + self.legacy.load_lora() 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