mirror of
https://github.com/invoke-ai/InvokeAI.git
synced 2026-02-19 09:54:24 -05:00
400 lines
14 KiB
Python
400 lines
14 KiB
Python
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 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
|
|
|
|
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 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 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):
|
|
# 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 <lora:lora_name:multiplier>
|
|
# 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:([^>]+)>")
|
|
lora_match = re.compile(r"withLora\(([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 clear_loras(self):
|
|
self.loras_to_load = {}
|
|
self.legacy.clear_loras()
|
|
|
|
def __del__(self):
|
|
del self.loras_to_load
|