ruizhaocv's picture
Upload 17 files
7d421db
raw
history blame
48.5 kB
import json
import math
from itertools import groupby
import os
from typing import Callable, Dict, List, Optional, Set, Tuple, Type, Union
import numpy as np
import PIL
import torch
import torch.nn as nn
import torch.nn.functional as F
try:
from safetensors.torch import safe_open
from safetensors.torch import save_file as safe_save
safetensors_available = True
except ImportError:
from .safe_open import safe_open
def safe_save(
tensors: Dict[str, torch.Tensor],
filename: str,
metadata: Optional[Dict[str, str]] = None,
) -> None:
raise EnvironmentError(
"Saving safetensors requires the safetensors library. Please install with pip or similar."
)
safetensors_available = False
class LoraInjectedLinear(nn.Module):
def __init__(
self, in_features, out_features, bias=False, r=4, dropout_p=0.1, scale=1.0
):
super().__init__()
if r > min(in_features, out_features):
#raise ValueError(
# f"LoRA rank {r} must be less or equal than {min(in_features, out_features)}"
#)
print(f"LoRA rank {r} is too large. setting to: {min(in_features, out_features)}")
r = min(in_features, out_features)
self.r = r
self.linear = nn.Linear(in_features, out_features, bias)
self.lora_down = nn.Linear(in_features, r, bias=False)
self.dropout = nn.Dropout(dropout_p)
self.lora_up = nn.Linear(r, out_features, bias=False)
self.scale = scale
self.selector = nn.Identity()
nn.init.normal_(self.lora_down.weight, std=1 / r)
nn.init.zeros_(self.lora_up.weight)
def forward(self, input):
return (
self.linear(input)
+ self.dropout(self.lora_up(self.selector(self.lora_down(input))))
* self.scale
)
def realize_as_lora(self):
return self.lora_up.weight.data * self.scale, self.lora_down.weight.data
def set_selector_from_diag(self, diag: torch.Tensor):
# diag is a 1D tensor of size (r,)
assert diag.shape == (self.r,)
self.selector = nn.Linear(self.r, self.r, bias=False)
self.selector.weight.data = torch.diag(diag)
self.selector.weight.data = self.selector.weight.data.to(
self.lora_up.weight.device
).to(self.lora_up.weight.dtype)
class MultiLoraInjectedLinear(nn.Module):
def __init__(
self, in_features, out_features, bias=False, r=4, dropout_p=0.1, lora_num=1, scales=[1.0]
):
super().__init__()
if r > min(in_features, out_features):
#raise ValueError(
# f"LoRA rank {r} must be less or equal than {min(in_features, out_features)}"
#)
print(f"LoRA rank {r} is too large. setting to: {min(in_features, out_features)}")
r = min(in_features, out_features)
self.r = r
self.linear = nn.Linear(in_features, out_features, bias)
for i in range(lora_num):
if i==0:
self.lora_down =[nn.Linear(in_features, r, bias=False)]
self.dropout = [nn.Dropout(dropout_p)]
self.lora_up = [nn.Linear(r, out_features, bias=False)]
self.scale = scales[i]
self.selector = [nn.Identity()]
else:
self.lora_down.append(nn.Linear(in_features, r, bias=False))
self.dropout.append( nn.Dropout(dropout_p))
self.lora_up.append( nn.Linear(r, out_features, bias=False))
self.scale.append(scales[i])
nn.init.normal_(self.lora_down.weight, std=1 / r)
nn.init.zeros_(self.lora_up.weight)
def forward(self, input):
return (
self.linear(input)
+ self.dropout(self.lora_up(self.selector(self.lora_down(input))))
* self.scale
)
def realize_as_lora(self):
return self.lora_up.weight.data * self.scale, self.lora_down.weight.data
def set_selector_from_diag(self, diag: torch.Tensor):
# diag is a 1D tensor of size (r,)
assert diag.shape == (self.r,)
self.selector = nn.Linear(self.r, self.r, bias=False)
self.selector.weight.data = torch.diag(diag)
self.selector.weight.data = self.selector.weight.data.to(
self.lora_up.weight.device
).to(self.lora_up.weight.dtype)
class LoraInjectedConv2d(nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size,
stride=1,
padding=0,
dilation=1,
groups: int = 1,
bias: bool = True,
r: int = 4,
dropout_p: float = 0.1,
scale: float = 1.0,
):
super().__init__()
if r > min(in_channels, out_channels):
print(f"LoRA rank {r} is too large. setting to: {min(in_channels, out_channels)}")
r = min(in_channels, out_channels)
self.r = r
self.conv = nn.Conv2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
bias=bias,
)
self.lora_down = nn.Conv2d(
in_channels=in_channels,
out_channels=r,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
bias=False,
)
self.dropout = nn.Dropout(dropout_p)
self.lora_up = nn.Conv2d(
in_channels=r,
out_channels=out_channels,
kernel_size=1,
stride=1,
padding=0,
bias=False,
)
self.selector = nn.Identity()
self.scale = scale
nn.init.normal_(self.lora_down.weight, std=1 / r)
nn.init.zeros_(self.lora_up.weight)
def forward(self, input):
return (
self.conv(input)
+ self.dropout(self.lora_up(self.selector(self.lora_down(input))))
* self.scale
)
def realize_as_lora(self):
return self.lora_up.weight.data * self.scale, self.lora_down.weight.data
def set_selector_from_diag(self, diag: torch.Tensor):
# diag is a 1D tensor of size (r,)
assert diag.shape == (self.r,)
self.selector = nn.Conv2d(
in_channels=self.r,
out_channels=self.r,
kernel_size=1,
stride=1,
padding=0,
bias=False,
)
self.selector.weight.data = torch.diag(diag)
# same device + dtype as lora_up
self.selector.weight.data = self.selector.weight.data.to(
self.lora_up.weight.device
).to(self.lora_up.weight.dtype)
class LoraInjectedConv3d(nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: (3, 1, 1),
padding: (1, 0, 0),
bias: bool = False,
r: int = 4,
dropout_p: float = 0,
scale: float = 1.0,
):
super().__init__()
if r > min(in_channels, out_channels):
print(f"LoRA rank {r} is too large. setting to: {min(in_channels, out_channels)}")
r = min(in_channels, out_channels)
self.r = r
self.kernel_size = kernel_size
self.padding = padding
self.conv = nn.Conv3d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
padding=padding,
)
self.lora_down = nn.Conv3d(
in_channels=in_channels,
out_channels=r,
kernel_size=kernel_size,
bias=False,
padding=padding
)
self.dropout = nn.Dropout(dropout_p)
self.lora_up = nn.Conv3d(
in_channels=r,
out_channels=out_channels,
kernel_size=1,
stride=1,
padding=0,
bias=False,
)
self.selector = nn.Identity()
self.scale = scale
nn.init.normal_(self.lora_down.weight, std=1 / r)
nn.init.zeros_(self.lora_up.weight)
def forward(self, input):
return (
self.conv(input)
+ self.dropout(self.lora_up(self.selector(self.lora_down(input))))
* self.scale
)
def realize_as_lora(self):
return self.lora_up.weight.data * self.scale, self.lora_down.weight.data
def set_selector_from_diag(self, diag: torch.Tensor):
# diag is a 1D tensor of size (r,)
assert diag.shape == (self.r,)
self.selector = nn.Conv3d(
in_channels=self.r,
out_channels=self.r,
kernel_size=1,
stride=1,
padding=0,
bias=False,
)
self.selector.weight.data = torch.diag(diag)
# same device + dtype as lora_up
self.selector.weight.data = self.selector.weight.data.to(
self.lora_up.weight.device
).to(self.lora_up.weight.dtype)
UNET_DEFAULT_TARGET_REPLACE = {"CrossAttention", "Attention", "GEGLU"}
UNET_EXTENDED_TARGET_REPLACE = {"ResnetBlock2D", "CrossAttention", "Attention", "GEGLU"}
TEXT_ENCODER_DEFAULT_TARGET_REPLACE = {"CLIPAttention"}
TEXT_ENCODER_EXTENDED_TARGET_REPLACE = {"CLIPAttention"}
DEFAULT_TARGET_REPLACE = UNET_DEFAULT_TARGET_REPLACE
EMBED_FLAG = "<embed>"
def _find_children(
model,
search_class: List[Type[nn.Module]] = [nn.Linear],
):
"""
Find all modules of a certain class (or union of classes).
Returns all matching modules, along with the parent of those moduless and the
names they are referenced by.
"""
# For each target find every linear_class module that isn't a child of a LoraInjectedLinear
for parent in model.modules():
for name, module in parent.named_children():
if any([isinstance(module, _class) for _class in search_class]):
yield parent, name, module
def _find_modules_v2(
model,
ancestor_class: Optional[Set[str]] = None,
search_class: List[Type[nn.Module]] = [nn.Linear],
exclude_children_of: Optional[List[Type[nn.Module]]] = None,
# [
# LoraInjectedLinear,
# LoraInjectedConv2d,
# LoraInjectedConv3d
# ],
):
"""
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 moduless 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 name, module in model.named_modules()
if module.__class__.__name__ in ancestor_class # and ('transformer_in' not in name)
)
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]):
continue_flag = True
if 'Transformer2DModel' in ancestor_class and ('attn1' in fullname or 'ff' in fullname):
continue_flag = False
if 'TransformerTemporalModel' in ancestor_class and ('attn1' in fullname or 'attn2' in fullname or 'ff' in fullname):
continue_flag = False
if continue_flag:
continue
# 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
if name in ['lora_up', 'dropout', 'lora_down']:
continue
# Otherwise, yield it
yield parent, name, module
def _find_modules_old(
model,
ancestor_class: Set[str] = DEFAULT_TARGET_REPLACE,
search_class: List[Type[nn.Module]] = [nn.Linear],
exclude_children_of: Optional[List[Type[nn.Module]]] = [LoraInjectedLinear],
):
ret = []
for _module in model.modules():
if _module.__class__.__name__ in ancestor_class:
for name, _child_module in _module.named_modules():
if _child_module.__class__ in search_class:
ret.append((_module, name, _child_module))
print(ret)
return ret
_find_modules = _find_modules_v2
def inject_trainable_lora(
model: nn.Module,
target_replace_module: Set[str] = DEFAULT_TARGET_REPLACE,
r: int = 4,
loras=None, # path to lora .pt
verbose: bool = False,
dropout_p: float = 0.0,
scale: float = 1.0,
):
"""
inject lora into model, and returns lora parameter groups.
"""
require_grad_params = []
names = []
if loras != None:
loras = torch.load(loras)
for _module, name, _child_module in _find_modules(
model, target_replace_module, search_class=[nn.Linear]
):
weight = _child_module.weight
bias = _child_module.bias
if verbose:
print("LoRA Injection : injecting lora into ", name)
print("LoRA Injection : weight shape", weight.shape)
_tmp = LoraInjectedLinear(
_child_module.in_features,
_child_module.out_features,
_child_module.bias is not None,
r=r,
dropout_p=dropout_p,
scale=scale,
)
_tmp.linear.weight = weight
if bias is not None:
_tmp.linear.bias = bias
# switch the module
_tmp.to(_child_module.weight.device).to(_child_module.weight.dtype)
_module._modules[name] = _tmp
require_grad_params.append(_module._modules[name].lora_up.parameters())
require_grad_params.append(_module._modules[name].lora_down.parameters())
if loras != None:
_module._modules[name].lora_up.weight = loras.pop(0)
_module._modules[name].lora_down.weight = loras.pop(0)
_module._modules[name].lora_up.weight.requires_grad = True
_module._modules[name].lora_down.weight.requires_grad = True
names.append(name)
return require_grad_params, names
def inject_trainable_lora_extended(
model: nn.Module,
target_replace_module: Set[str] = UNET_EXTENDED_TARGET_REPLACE,
r: int = 4,
loras=None, # path to lora .pt
dropout_p: float = 0.0,
scale: float = 1.0,
):
"""
inject lora into model, and returns lora parameter groups.
"""
require_grad_params = []
names = []
if loras != None:
loras = torch.load(loras)
if True:
for target_replace_module_i in target_replace_module:
for _module, name, _child_module in _find_modules(
model, [target_replace_module_i], search_class=[nn.Linear, nn.Conv2d, nn.Conv3d]
):
# if name == 'to_q':
# continue
if _child_module.__class__ == nn.Linear:
weight = _child_module.weight
bias = _child_module.bias
_tmp = LoraInjectedLinear(
_child_module.in_features,
_child_module.out_features,
_child_module.bias is not None,
r=r,
dropout_p=dropout_p,
scale=scale,
)
_tmp.linear.weight = weight
if bias is not None:
_tmp.linear.bias = bias
elif _child_module.__class__ == nn.Conv2d:
weight = _child_module.weight
bias = _child_module.bias
_tmp = LoraInjectedConv2d(
_child_module.in_channels,
_child_module.out_channels,
_child_module.kernel_size,
_child_module.stride,
_child_module.padding,
_child_module.dilation,
_child_module.groups,
_child_module.bias is not None,
r=r,
dropout_p=dropout_p,
scale=scale,
)
_tmp.conv.weight = weight
if bias is not None:
_tmp.conv.bias = bias
elif _child_module.__class__ == nn.Conv3d:
weight = _child_module.weight
bias = _child_module.bias
_tmp = LoraInjectedConv3d(
_child_module.in_channels,
_child_module.out_channels,
bias=_child_module.bias is not None,
kernel_size=_child_module.kernel_size,
padding=_child_module.padding,
r=r,
dropout_p=dropout_p,
scale=scale,
)
_tmp.conv.weight = weight
if bias is not None:
_tmp.conv.bias = bias
# switch the module
_tmp.to(_child_module.weight.device).to(_child_module.weight.dtype)
if bias is not None:
_tmp.to(_child_module.bias.device).to(_child_module.bias.dtype)
_module._modules[name] = _tmp
require_grad_params.append(_module._modules[name].lora_up.parameters())
require_grad_params.append(_module._modules[name].lora_down.parameters())
if loras != None:
_module._modules[name].lora_up.weight = loras.pop(0)
_module._modules[name].lora_down.weight = loras.pop(0)
_module._modules[name].lora_up.weight.requires_grad = True
_module._modules[name].lora_down.weight.requires_grad = True
names.append(name)
else:
for _module, name, _child_module in _find_modules(
model, target_replace_module, search_class=[nn.Linear, nn.Conv2d, nn.Conv3d]
):
if _child_module.__class__ == nn.Linear:
weight = _child_module.weight
bias = _child_module.bias
_tmp = LoraInjectedLinear(
_child_module.in_features,
_child_module.out_features,
_child_module.bias is not None,
r=r,
dropout_p=dropout_p,
scale=scale,
)
_tmp.linear.weight = weight
if bias is not None:
_tmp.linear.bias = bias
elif _child_module.__class__ == nn.Conv2d:
weight = _child_module.weight
bias = _child_module.bias
_tmp = LoraInjectedConv2d(
_child_module.in_channels,
_child_module.out_channels,
_child_module.kernel_size,
_child_module.stride,
_child_module.padding,
_child_module.dilation,
_child_module.groups,
_child_module.bias is not None,
r=r,
dropout_p=dropout_p,
scale=scale,
)
_tmp.conv.weight = weight
if bias is not None:
_tmp.conv.bias = bias
elif _child_module.__class__ == nn.Conv3d:
weight = _child_module.weight
bias = _child_module.bias
_tmp = LoraInjectedConv3d(
_child_module.in_channels,
_child_module.out_channels,
bias=_child_module.bias is not None,
kernel_size=_child_module.kernel_size,
padding=_child_module.padding,
r=r,
dropout_p=dropout_p,
scale=scale,
)
_tmp.conv.weight = weight
if bias is not None:
_tmp.conv.bias = bias
# switch the module
_tmp.to(_child_module.weight.device).to(_child_module.weight.dtype)
if bias is not None:
_tmp.to(_child_module.bias.device).to(_child_module.bias.dtype)
_module._modules[name] = _tmp
require_grad_params.append(_module._modules[name].lora_up.parameters())
require_grad_params.append(_module._modules[name].lora_down.parameters())
if loras != None:
_module._modules[name].lora_up.weight = loras.pop(0)
_module._modules[name].lora_down.weight = loras.pop(0)
_module._modules[name].lora_up.weight.requires_grad = True
_module._modules[name].lora_down.weight.requires_grad = True
names.append(name)
return require_grad_params, names
def inject_inferable_lora(
model,
lora_path='',
unet_replace_modules=["UNet3DConditionModel"],
text_encoder_replace_modules=["CLIPEncoderLayer"],
is_extended=False,
r=16
):
from transformers.models.clip import CLIPTextModel
from diffusers import UNet3DConditionModel
def is_text_model(f): return 'text_encoder' in f and isinstance(model.text_encoder, CLIPTextModel)
def is_unet(f): return 'unet' in f and model.unet.__class__.__name__ == "UNet3DConditionModel"
if os.path.exists(lora_path):
try:
for f in os.listdir(lora_path):
if f.endswith('.pt'):
lora_file = os.path.join(lora_path, f)
if is_text_model(f):
monkeypatch_or_replace_lora(
model.text_encoder,
torch.load(lora_file),
target_replace_module=text_encoder_replace_modules,
r=r
)
print("Successfully loaded Text Encoder LoRa.")
continue
if is_unet(f):
monkeypatch_or_replace_lora_extended(
model.unet,
torch.load(lora_file),
target_replace_module=unet_replace_modules,
r=r
)
print("Successfully loaded UNET LoRa.")
continue
print("Found a .pt file, but doesn't have the correct name format. (unet.pt, text_encoder.pt)")
except Exception as e:
print(e)
print("Couldn't inject LoRA's due to an error.")
def extract_lora_ups_down(model, target_replace_module=DEFAULT_TARGET_REPLACE):
loras = []
for target_replace_module_i in target_replace_module:
for _m, _n, _child_module in _find_modules(
model,
[target_replace_module_i],
search_class=[LoraInjectedLinear, LoraInjectedConv2d, LoraInjectedConv3d],
):
loras.append((_child_module.lora_up, _child_module.lora_down))
if len(loras) == 0:
raise ValueError("No lora injected.")
return loras
def extract_lora_child_module(model, target_replace_module=DEFAULT_TARGET_REPLACE):
loras = []
for target_replace_module_i in target_replace_module:
for _m, _n, _child_module in _find_modules(
model,
[target_replace_module_i],
search_class=[LoraInjectedLinear, LoraInjectedConv2d, LoraInjectedConv3d],
):
loras.append(_child_module)
if len(loras) == 0:
raise ValueError("No lora injected.")
return loras
def extract_lora_as_tensor(
model, target_replace_module=DEFAULT_TARGET_REPLACE, as_fp16=True
):
loras = []
for _m, _n, _child_module in _find_modules(
model,
target_replace_module,
search_class=[LoraInjectedLinear, LoraInjectedConv2d, LoraInjectedConv3d],
):
up, down = _child_module.realize_as_lora()
if as_fp16:
up = up.to(torch.float16)
down = down.to(torch.float16)
loras.append((up, down))
if len(loras) == 0:
raise ValueError("No lora injected.")
return loras
def save_lora_weight(
model,
path="./lora.pt",
target_replace_module=DEFAULT_TARGET_REPLACE,
flag=None
):
weights = []
for _up, _down in extract_lora_ups_down(
model, target_replace_module=target_replace_module
):
weights.append(_up.weight.to("cpu").to(torch.float32))
weights.append(_down.weight.to("cpu").to(torch.float32))
if not flag:
torch.save(weights, path)
else:
weights_new=[]
for i in range(0, len(weights), 4):
subset = weights[i+(flag-1)*2:i+(flag-1)*2+2]
weights_new.extend(subset)
torch.save(weights_new, path)
def save_lora_as_json(model, path="./lora.json"):
weights = []
for _up, _down in extract_lora_ups_down(model):
weights.append(_up.weight.detach().cpu().numpy().tolist())
weights.append(_down.weight.detach().cpu().numpy().tolist())
import json
with open(path, "w") as f:
json.dump(weights, f)
def save_safeloras_with_embeds(
modelmap: Dict[str, Tuple[nn.Module, Set[str]]] = {},
embeds: Dict[str, torch.Tensor] = {},
outpath="./lora.safetensors",
):
"""
Saves the Lora from multiple modules in a single safetensor file.
modelmap is a dictionary of {
"module name": (module, target_replace_module)
}
"""
weights = {}
metadata = {}
for name, (model, target_replace_module) in modelmap.items():
metadata[name] = json.dumps(list(target_replace_module))
for i, (_up, _down) in enumerate(
extract_lora_as_tensor(model, target_replace_module)
):
rank = _down.shape[0]
metadata[f"{name}:{i}:rank"] = str(rank)
weights[f"{name}:{i}:up"] = _up
weights[f"{name}:{i}:down"] = _down
for token, tensor in embeds.items():
metadata[token] = EMBED_FLAG
weights[token] = tensor
print(f"Saving weights to {outpath}")
safe_save(weights, outpath, metadata)
def save_safeloras(
modelmap: Dict[str, Tuple[nn.Module, Set[str]]] = {},
outpath="./lora.safetensors",
):
return save_safeloras_with_embeds(modelmap=modelmap, outpath=outpath)
def convert_loras_to_safeloras_with_embeds(
modelmap: Dict[str, Tuple[str, Set[str], int]] = {},
embeds: Dict[str, torch.Tensor] = {},
outpath="./lora.safetensors",
):
"""
Converts the Lora from multiple pytorch .pt files into a single safetensor file.
modelmap is a dictionary of {
"module name": (pytorch_model_path, target_replace_module, rank)
}
"""
weights = {}
metadata = {}
for name, (path, target_replace_module, r) in modelmap.items():
metadata[name] = json.dumps(list(target_replace_module))
lora = torch.load(path)
for i, weight in enumerate(lora):
is_up = i % 2 == 0
i = i // 2
if is_up:
metadata[f"{name}:{i}:rank"] = str(r)
weights[f"{name}:{i}:up"] = weight
else:
weights[f"{name}:{i}:down"] = weight
for token, tensor in embeds.items():
metadata[token] = EMBED_FLAG
weights[token] = tensor
print(f"Saving weights to {outpath}")
safe_save(weights, outpath, metadata)
def convert_loras_to_safeloras(
modelmap: Dict[str, Tuple[str, Set[str], int]] = {},
outpath="./lora.safetensors",
):
convert_loras_to_safeloras_with_embeds(modelmap=modelmap, outpath=outpath)
def parse_safeloras(
safeloras,
) -> Dict[str, Tuple[List[nn.parameter.Parameter], List[int], List[str]]]:
"""
Converts a loaded safetensor file that contains a set of module Loras
into Parameters and other information
Output is a dictionary of {
"module name": (
[list of weights],
[list of ranks],
target_replacement_modules
)
}
"""
loras = {}
metadata = safeloras.metadata()
get_name = lambda k: k.split(":")[0]
keys = list(safeloras.keys())
keys.sort(key=get_name)
for name, module_keys in groupby(keys, get_name):
info = metadata.get(name)
if not info:
raise ValueError(
f"Tensor {name} has no metadata - is this a Lora safetensor?"
)
# Skip Textual Inversion embeds
if info == EMBED_FLAG:
continue
# Handle Loras
# Extract the targets
target = json.loads(info)
# Build the result lists - Python needs us to preallocate lists to insert into them
module_keys = list(module_keys)
ranks = [4] * (len(module_keys) // 2)
weights = [None] * len(module_keys)
for key in module_keys:
# Split the model name and index out of the key
_, idx, direction = key.split(":")
idx = int(idx)
# Add the rank
ranks[idx] = int(metadata[f"{name}:{idx}:rank"])
# Insert the weight into the list
idx = idx * 2 + (1 if direction == "down" else 0)
weights[idx] = nn.parameter.Parameter(safeloras.get_tensor(key))
loras[name] = (weights, ranks, target)
return loras
def parse_safeloras_embeds(
safeloras,
) -> Dict[str, torch.Tensor]:
"""
Converts a loaded safetensor file that contains Textual Inversion embeds into
a dictionary of embed_token: Tensor
"""
embeds = {}
metadata = safeloras.metadata()
for key in safeloras.keys():
# Only handle Textual Inversion embeds
meta = metadata.get(key)
if not meta or meta != EMBED_FLAG:
continue
embeds[key] = safeloras.get_tensor(key)
return embeds
def load_safeloras(path, device="cpu"):
safeloras = safe_open(path, framework="pt", device=device)
return parse_safeloras(safeloras)
def load_safeloras_embeds(path, device="cpu"):
safeloras = safe_open(path, framework="pt", device=device)
return parse_safeloras_embeds(safeloras)
def load_safeloras_both(path, device="cpu"):
safeloras = safe_open(path, framework="pt", device=device)
return parse_safeloras(safeloras), parse_safeloras_embeds(safeloras)
def collapse_lora(model, alpha=1.0):
for _module, name, _child_module in _find_modules(
model,
UNET_EXTENDED_TARGET_REPLACE | TEXT_ENCODER_EXTENDED_TARGET_REPLACE,
search_class=[LoraInjectedLinear, LoraInjectedConv2d, LoraInjectedConv3d],
):
if isinstance(_child_module, LoraInjectedLinear):
print("Collapsing Lin Lora in", name)
_child_module.linear.weight = nn.Parameter(
_child_module.linear.weight.data
+ alpha
* (
_child_module.lora_up.weight.data
@ _child_module.lora_down.weight.data
)
.type(_child_module.linear.weight.dtype)
.to(_child_module.linear.weight.device)
)
else:
print("Collapsing Conv Lora in", name)
_child_module.conv.weight = nn.Parameter(
_child_module.conv.weight.data
+ alpha
* (
_child_module.lora_up.weight.data.flatten(start_dim=1)
@ _child_module.lora_down.weight.data.flatten(start_dim=1)
)
.reshape(_child_module.conv.weight.data.shape)
.type(_child_module.conv.weight.dtype)
.to(_child_module.conv.weight.device)
)
def monkeypatch_or_replace_lora(
model,
loras,
target_replace_module=DEFAULT_TARGET_REPLACE,
r: Union[int, List[int]] = 4,
):
for _module, name, _child_module in _find_modules(
model, target_replace_module, search_class=[nn.Linear, LoraInjectedLinear]
):
_source = (
_child_module.linear
if isinstance(_child_module, LoraInjectedLinear)
else _child_module
)
weight = _source.weight
bias = _source.bias
_tmp = LoraInjectedLinear(
_source.in_features,
_source.out_features,
_source.bias is not None,
r=r.pop(0) if isinstance(r, list) else r,
)
_tmp.linear.weight = weight
if bias is not None:
_tmp.linear.bias = bias
# switch the module
_module._modules[name] = _tmp
up_weight = loras.pop(0)
down_weight = loras.pop(0)
_module._modules[name].lora_up.weight = nn.Parameter(
up_weight.type(weight.dtype)
)
_module._modules[name].lora_down.weight = nn.Parameter(
down_weight.type(weight.dtype)
)
_module._modules[name].to(weight.device)
def monkeypatch_or_replace_lora_extended(
model,
loras,
target_replace_module=DEFAULT_TARGET_REPLACE,
r: Union[int, List[int]] = 4,
):
for _module, name, _child_module in _find_modules(
model,
target_replace_module,
search_class=[
nn.Linear,
nn.Conv2d,
nn.Conv3d,
LoraInjectedLinear,
LoraInjectedConv2d,
LoraInjectedConv3d,
],
):
if (_child_module.__class__ == nn.Linear) or (
_child_module.__class__ == LoraInjectedLinear
):
if len(loras[0].shape) != 2:
continue
_source = (
_child_module.linear
if isinstance(_child_module, LoraInjectedLinear)
else _child_module
)
weight = _source.weight
bias = _source.bias
_tmp = LoraInjectedLinear(
_source.in_features,
_source.out_features,
_source.bias is not None,
r=r.pop(0) if isinstance(r, list) else r,
)
_tmp.linear.weight = weight
if bias is not None:
_tmp.linear.bias = bias
elif (_child_module.__class__ == nn.Conv2d) or (
_child_module.__class__ == LoraInjectedConv2d
):
if len(loras[0].shape) != 4:
continue
_source = (
_child_module.conv
if isinstance(_child_module, LoraInjectedConv2d)
else _child_module
)
weight = _source.weight
bias = _source.bias
_tmp = LoraInjectedConv2d(
_source.in_channels,
_source.out_channels,
_source.kernel_size,
_source.stride,
_source.padding,
_source.dilation,
_source.groups,
_source.bias is not None,
r=r.pop(0) if isinstance(r, list) else r,
)
_tmp.conv.weight = weight
if bias is not None:
_tmp.conv.bias = bias
elif _child_module.__class__ == nn.Conv3d or(
_child_module.__class__ == LoraInjectedConv3d
):
if len(loras[0].shape) != 5:
continue
_source = (
_child_module.conv
if isinstance(_child_module, LoraInjectedConv3d)
else _child_module
)
weight = _source.weight
bias = _source.bias
_tmp = LoraInjectedConv3d(
_source.in_channels,
_source.out_channels,
bias=_source.bias is not None,
kernel_size=_source.kernel_size,
padding=_source.padding,
r=r.pop(0) if isinstance(r, list) else r,
)
_tmp.conv.weight = weight
if bias is not None:
_tmp.conv.bias = bias
# switch the module
_module._modules[name] = _tmp
up_weight = loras.pop(0)
down_weight = loras.pop(0)
_module._modules[name].lora_up.weight = nn.Parameter(
up_weight.type(weight.dtype)
)
_module._modules[name].lora_down.weight = nn.Parameter(
down_weight.type(weight.dtype)
)
_module._modules[name].to(weight.device)
def monkeypatch_or_replace_safeloras(models, safeloras):
loras = parse_safeloras(safeloras)
for name, (lora, ranks, target) in loras.items():
model = getattr(models, name, None)
if not model:
print(f"No model provided for {name}, contained in Lora")
continue
monkeypatch_or_replace_lora_extended(model, lora, target, ranks)
def monkeypatch_remove_lora(model):
for _module, name, _child_module in _find_modules(
model, search_class=[LoraInjectedLinear, LoraInjectedConv2d, LoraInjectedConv3d]
):
if isinstance(_child_module, LoraInjectedLinear):
_source = _child_module.linear
weight, bias = _source.weight, _source.bias
_tmp = nn.Linear(
_source.in_features, _source.out_features, bias is not None
)
_tmp.weight = weight
if bias is not None:
_tmp.bias = bias
else:
_source = _child_module.conv
weight, bias = _source.weight, _source.bias
if isinstance(_source, nn.Conv2d):
_tmp = nn.Conv2d(
in_channels=_source.in_channels,
out_channels=_source.out_channels,
kernel_size=_source.kernel_size,
stride=_source.stride,
padding=_source.padding,
dilation=_source.dilation,
groups=_source.groups,
bias=bias is not None,
)
_tmp.weight = weight
if bias is not None:
_tmp.bias = bias
if isinstance(_source, nn.Conv3d):
_tmp = nn.Conv3d(
_source.in_channels,
_source.out_channels,
bias=_source.bias is not None,
kernel_size=_source.kernel_size,
padding=_source.padding,
)
_tmp.weight = weight
if bias is not None:
_tmp.bias = bias
_module._modules[name] = _tmp
def monkeypatch_add_lora(
model,
loras,
target_replace_module=DEFAULT_TARGET_REPLACE,
alpha: float = 1.0,
beta: float = 1.0,
):
for _module, name, _child_module in _find_modules(
model, target_replace_module, search_class=[LoraInjectedLinear]
):
weight = _child_module.linear.weight
up_weight = loras.pop(0)
down_weight = loras.pop(0)
_module._modules[name].lora_up.weight = nn.Parameter(
up_weight.type(weight.dtype).to(weight.device) * alpha
+ _module._modules[name].lora_up.weight.to(weight.device) * beta
)
_module._modules[name].lora_down.weight = nn.Parameter(
down_weight.type(weight.dtype).to(weight.device) * alpha
+ _module._modules[name].lora_down.weight.to(weight.device) * beta
)
_module._modules[name].to(weight.device)
def tune_lora_scale(model, alpha: float = 1.0):
for _module in model.modules():
if _module.__class__.__name__ in ["LoraInjectedLinear", "LoraInjectedConv2d", "LoraInjectedConv3d"]:
_module.scale = alpha
def set_lora_diag(model, diag: torch.Tensor):
for _module in model.modules():
if _module.__class__.__name__ in ["LoraInjectedLinear", "LoraInjectedConv2d", "LoraInjectedConv3d"]:
_module.set_selector_from_diag(diag)
def _text_lora_path(path: str) -> str:
assert path.endswith(".pt"), "Only .pt files are supported"
return ".".join(path.split(".")[:-1] + ["text_encoder", "pt"])
def _ti_lora_path(path: str) -> str:
assert path.endswith(".pt"), "Only .pt files are supported"
return ".".join(path.split(".")[:-1] + ["ti", "pt"])
def apply_learned_embed_in_clip(
learned_embeds,
text_encoder,
tokenizer,
token: Optional[Union[str, List[str]]] = None,
idempotent=False,
):
if isinstance(token, str):
trained_tokens = [token]
elif isinstance(token, list):
assert len(learned_embeds.keys()) == len(
token
), "The number of tokens and the number of embeds should be the same"
trained_tokens = token
else:
trained_tokens = list(learned_embeds.keys())
for token in trained_tokens:
print(token)
embeds = learned_embeds[token]
# cast to dtype of text_encoder
dtype = text_encoder.get_input_embeddings().weight.dtype
num_added_tokens = tokenizer.add_tokens(token)
i = 1
if not idempotent:
while num_added_tokens == 0:
print(f"The tokenizer already contains the token {token}.")
token = f"{token[:-1]}-{i}>"
print(f"Attempting to add the token {token}.")
num_added_tokens = tokenizer.add_tokens(token)
i += 1
elif num_added_tokens == 0 and idempotent:
print(f"The tokenizer already contains the token {token}.")
print(f"Replacing {token} embedding.")
# resize the token embeddings
text_encoder.resize_token_embeddings(len(tokenizer))
# get the id for the token and assign the embeds
token_id = tokenizer.convert_tokens_to_ids(token)
text_encoder.get_input_embeddings().weight.data[token_id] = embeds
return token
def load_learned_embed_in_clip(
learned_embeds_path,
text_encoder,
tokenizer,
token: Optional[Union[str, List[str]]] = None,
idempotent=False,
):
learned_embeds = torch.load(learned_embeds_path)
apply_learned_embed_in_clip(
learned_embeds, text_encoder, tokenizer, token, idempotent
)
def patch_pipe(
pipe,
maybe_unet_path,
token: Optional[str] = None,
r: int = 4,
patch_unet=True,
patch_text=True,
patch_ti=True,
idempotent_token=True,
unet_target_replace_module=DEFAULT_TARGET_REPLACE,
text_target_replace_module=TEXT_ENCODER_DEFAULT_TARGET_REPLACE,
):
if maybe_unet_path.endswith(".pt"):
# torch format
if maybe_unet_path.endswith(".ti.pt"):
unet_path = maybe_unet_path[:-6] + ".pt"
elif maybe_unet_path.endswith(".text_encoder.pt"):
unet_path = maybe_unet_path[:-16] + ".pt"
else:
unet_path = maybe_unet_path
ti_path = _ti_lora_path(unet_path)
text_path = _text_lora_path(unet_path)
if patch_unet:
print("LoRA : Patching Unet")
monkeypatch_or_replace_lora(
pipe.unet,
torch.load(unet_path),
r=r,
target_replace_module=unet_target_replace_module,
)
if patch_text:
print("LoRA : Patching text encoder")
monkeypatch_or_replace_lora(
pipe.text_encoder,
torch.load(text_path),
target_replace_module=text_target_replace_module,
r=r,
)
if patch_ti:
print("LoRA : Patching token input")
token = load_learned_embed_in_clip(
ti_path,
pipe.text_encoder,
pipe.tokenizer,
token=token,
idempotent=idempotent_token,
)
elif maybe_unet_path.endswith(".safetensors"):
safeloras = safe_open(maybe_unet_path, framework="pt", device="cpu")
monkeypatch_or_replace_safeloras(pipe, safeloras)
tok_dict = parse_safeloras_embeds(safeloras)
if patch_ti:
apply_learned_embed_in_clip(
tok_dict,
pipe.text_encoder,
pipe.tokenizer,
token=token,
idempotent=idempotent_token,
)
return tok_dict
def train_patch_pipe(pipe, patch_unet, patch_text):
if patch_unet:
print("LoRA : Patching Unet")
collapse_lora(pipe.unet)
monkeypatch_remove_lora(pipe.unet)
if patch_text:
print("LoRA : Patching text encoder")
collapse_lora(pipe.text_encoder)
monkeypatch_remove_lora(pipe.text_encoder)
@torch.no_grad()
def inspect_lora(model):
moved = {}
for name, _module in model.named_modules():
if _module.__class__.__name__ in ["LoraInjectedLinear", "LoraInjectedConv2d", "LoraInjectedConv3d"]:
ups = _module.lora_up.weight.data.clone()
downs = _module.lora_down.weight.data.clone()
wght: torch.Tensor = ups.flatten(1) @ downs.flatten(1)
dist = wght.flatten().abs().mean().item()
if name in moved:
moved[name].append(dist)
else:
moved[name] = [dist]
return moved
def save_all(
unet,
text_encoder,
save_path,
placeholder_token_ids=None,
placeholder_tokens=None,
save_lora=True,
save_ti=True,
target_replace_module_text=TEXT_ENCODER_DEFAULT_TARGET_REPLACE,
target_replace_module_unet=DEFAULT_TARGET_REPLACE,
safe_form=True,
):
if not safe_form:
# save ti
if save_ti:
ti_path = _ti_lora_path(save_path)
learned_embeds_dict = {}
for tok, tok_id in zip(placeholder_tokens, placeholder_token_ids):
learned_embeds = text_encoder.get_input_embeddings().weight[tok_id]
print(
f"Current Learned Embeddings for {tok}:, id {tok_id} ",
learned_embeds[:4],
)
learned_embeds_dict[tok] = learned_embeds.detach().cpu()
torch.save(learned_embeds_dict, ti_path)
print("Ti saved to ", ti_path)
# save text encoder
if save_lora:
save_lora_weight(
unet, save_path, target_replace_module=target_replace_module_unet
)
print("Unet saved to ", save_path)
save_lora_weight(
text_encoder,
_text_lora_path(save_path),
target_replace_module=target_replace_module_text,
)
print("Text Encoder saved to ", _text_lora_path(save_path))
else:
assert save_path.endswith(
".safetensors"
), f"Save path : {save_path} should end with .safetensors"
loras = {}
embeds = {}
if save_lora:
loras["unet"] = (unet, target_replace_module_unet)
loras["text_encoder"] = (text_encoder, target_replace_module_text)
if save_ti:
for tok, tok_id in zip(placeholder_tokens, placeholder_token_ids):
learned_embeds = text_encoder.get_input_embeddings().weight[tok_id]
print(
f"Current Learned Embeddings for {tok}:, id {tok_id} ",
learned_embeds[:4],
)
embeds[tok] = learned_embeds.detach().cpu()
save_safeloras_with_embeds(loras, embeds, save_path)