mirror of
https://github.com/kohya-ss/sd-scripts.git
synced 2026-04-08 22:35:09 +00:00
779 lines
32 KiB
Python
779 lines
32 KiB
Python
# LoRA network module
|
|
# reference:
|
|
# https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
|
|
# https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py
|
|
|
|
import math
|
|
import os
|
|
from typing import List, Tuple, Union
|
|
import numpy as np
|
|
import torch
|
|
import re
|
|
|
|
from library import train_util
|
|
|
|
RE_UPDOWN = re.compile(r"(up|down)_blocks_(\d+)_(resnets|upsamplers|downsamplers|attentions)_(\d+)_")
|
|
|
|
|
|
class LoRAModule(torch.nn.Module):
|
|
"""
|
|
replaces forward method of the original Linear, instead of replacing the original Linear module.
|
|
"""
|
|
|
|
def __init__(self, lora_name, org_module: torch.nn.Module, multiplier=1.0, lora_dim=4, alpha=1):
|
|
"""if alpha == 0 or None, alpha is rank (no scaling)."""
|
|
super().__init__()
|
|
self.lora_name = lora_name
|
|
|
|
if org_module.__class__.__name__ == "Conv2d":
|
|
in_dim = org_module.in_channels
|
|
out_dim = org_module.out_channels
|
|
else:
|
|
in_dim = org_module.in_features
|
|
out_dim = org_module.out_features
|
|
|
|
# if limit_rank:
|
|
# self.lora_dim = min(lora_dim, in_dim, out_dim)
|
|
# if self.lora_dim != lora_dim:
|
|
# print(f"{lora_name} dim (rank) is changed to: {self.lora_dim}")
|
|
# else:
|
|
self.lora_dim = lora_dim
|
|
|
|
if org_module.__class__.__name__ == "Conv2d":
|
|
kernel_size = org_module.kernel_size
|
|
stride = org_module.stride
|
|
padding = org_module.padding
|
|
self.lora_down = torch.nn.Conv2d(in_dim, self.lora_dim, kernel_size, stride, padding, bias=False)
|
|
self.lora_up = torch.nn.Conv2d(self.lora_dim, out_dim, (1, 1), (1, 1), bias=False)
|
|
else:
|
|
self.lora_down = torch.nn.Linear(in_dim, self.lora_dim, bias=False)
|
|
self.lora_up = torch.nn.Linear(self.lora_dim, out_dim, bias=False)
|
|
|
|
if type(alpha) == torch.Tensor:
|
|
alpha = alpha.detach().float().numpy() # without casting, bf16 causes error
|
|
alpha = self.lora_dim if alpha is None or alpha == 0 else alpha
|
|
self.scale = alpha / self.lora_dim
|
|
self.register_buffer("alpha", torch.tensor(alpha)) # 定数として扱える
|
|
|
|
# same as microsoft's
|
|
torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
|
|
torch.nn.init.zeros_(self.lora_up.weight)
|
|
|
|
self.multiplier = multiplier
|
|
self.org_module = org_module # remove in applying
|
|
self.region = None
|
|
self.region_mask = None
|
|
|
|
def apply_to(self):
|
|
self.org_forward = self.org_module.forward
|
|
self.org_module.forward = self.forward
|
|
del self.org_module
|
|
|
|
def merge_to(self, sd, dtype, device):
|
|
# get up/down weight
|
|
up_weight = sd["lora_up.weight"].to(torch.float).to(device)
|
|
down_weight = sd["lora_down.weight"].to(torch.float).to(device)
|
|
|
|
# extract weight from org_module
|
|
org_sd = self.org_module.state_dict()
|
|
weight = org_sd["weight"].to(torch.float)
|
|
|
|
# merge weight
|
|
if len(weight.size()) == 2:
|
|
# linear
|
|
weight = weight + self.multiplier * (up_weight @ down_weight) * self.scale
|
|
elif down_weight.size()[2:4] == (1, 1):
|
|
# conv2d 1x1
|
|
weight = (
|
|
weight
|
|
+ self.multiplier
|
|
* (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3)
|
|
* self.scale
|
|
)
|
|
else:
|
|
# conv2d 3x3
|
|
conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
|
|
# print(conved.size(), weight.size(), module.stride, module.padding)
|
|
weight = weight + self.multiplier * conved * self.scale
|
|
|
|
# set weight to org_module
|
|
org_sd["weight"] = weight.to(dtype)
|
|
self.org_module.load_state_dict(org_sd)
|
|
|
|
def set_region(self, region):
|
|
self.region = region
|
|
self.region_mask = None
|
|
|
|
def forward(self, x):
|
|
if self.region is None:
|
|
return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
|
|
|
|
# regional LoRA FIXME same as additional-network extension
|
|
if x.size()[1] % 77 == 0:
|
|
# print(f"LoRA for context: {self.lora_name}")
|
|
self.region = None
|
|
return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
|
|
|
|
# calculate region mask first time
|
|
if self.region_mask is None:
|
|
if len(x.size()) == 4:
|
|
h, w = x.size()[2:4]
|
|
else:
|
|
seq_len = x.size()[1]
|
|
ratio = math.sqrt((self.region.size()[0] * self.region.size()[1]) / seq_len)
|
|
h = int(self.region.size()[0] / ratio + 0.5)
|
|
w = seq_len // h
|
|
|
|
r = self.region.to(x.device)
|
|
if r.dtype == torch.bfloat16:
|
|
r = r.to(torch.float)
|
|
r = r.unsqueeze(0).unsqueeze(1)
|
|
# print(self.lora_name, self.region.size(), x.size(), r.size(), h, w)
|
|
r = torch.nn.functional.interpolate(r, (h, w), mode="bilinear")
|
|
r = r.to(x.dtype)
|
|
|
|
if len(x.size()) == 3:
|
|
r = torch.reshape(r, (1, x.size()[1], -1))
|
|
|
|
self.region_mask = r
|
|
|
|
return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale * self.region_mask
|
|
|
|
|
|
def create_network(multiplier, network_dim, network_alpha, vae, text_encoder, unet, **kwargs):
|
|
if network_dim is None:
|
|
network_dim = 4 # default
|
|
if network_alpha is None:
|
|
network_alpha = 1.0
|
|
|
|
# extract dim/alpha for conv2d, and block dim
|
|
conv_dim = kwargs.get("conv_dim", None)
|
|
conv_alpha = kwargs.get("conv_alpha", None)
|
|
if conv_dim is not None:
|
|
conv_dim = int(conv_dim)
|
|
if conv_alpha is None:
|
|
conv_alpha = 1.0
|
|
else:
|
|
conv_alpha = float(conv_alpha)
|
|
|
|
# block dim/alpha/lr
|
|
block_dims = kwargs.get("block_dims", None)
|
|
down_lr_weight = kwargs.get("down_lr_weight", None)
|
|
mid_lr_weight = kwargs.get("mid_lr_weight", None)
|
|
up_lr_weight = kwargs.get("up_lr_weight", None)
|
|
|
|
# 以上のいずれかに指定があればblockごとのdim(rank)を有効にする
|
|
if block_dims is not None or down_lr_weight is not None or mid_lr_weight is not None or up_lr_weight is not None:
|
|
block_alphas = kwargs.get("block_alphas", None)
|
|
conv_block_dims = kwargs.get("conv_block_dims", None)
|
|
conv_block_alphas = kwargs.get("conv_block_alphas", None)
|
|
|
|
block_dims, block_alphas, conv_block_dims, conv_block_alphas = get_block_dims_and_alphas(
|
|
block_dims, block_alphas, network_dim, network_alpha, conv_block_dims, conv_block_alphas, conv_dim, conv_alpha
|
|
)
|
|
|
|
# extract learning rate weight for each block
|
|
if down_lr_weight is not None:
|
|
# if some parameters are not set, use zero
|
|
if "," in down_lr_weight:
|
|
down_lr_weight = [(float(s) if s else 0.0) for s in down_lr_weight.split(",")]
|
|
|
|
if mid_lr_weight is not None:
|
|
mid_lr_weight = float(mid_lr_weight)
|
|
|
|
if up_lr_weight is not None:
|
|
if "," in up_lr_weight:
|
|
up_lr_weight = [(float(s) if s else 0.0) for s in up_lr_weight.split(",")]
|
|
|
|
down_lr_weight, mid_lr_weight, up_lr_weight = get_block_lr_weight(
|
|
down_lr_weight, mid_lr_weight, up_lr_weight, float(kwargs.get("block_lr_zero_threshold", 0.0))
|
|
)
|
|
|
|
# remove block dim/alpha without learning rate
|
|
block_dims, block_alphas, conv_block_dims, conv_block_alphas = remove_block_dims_and_alphas(
|
|
block_dims, block_alphas, conv_block_dims, conv_block_alphas, down_lr_weight, mid_lr_weight, up_lr_weight
|
|
)
|
|
|
|
else:
|
|
block_alphas = None
|
|
conv_block_dims = None
|
|
conv_block_alphas = None
|
|
|
|
# すごく引数が多いな ( ^ω^)･･･
|
|
network = LoRANetwork(
|
|
text_encoder,
|
|
unet,
|
|
multiplier=multiplier,
|
|
lora_dim=network_dim,
|
|
alpha=network_alpha,
|
|
conv_lora_dim=conv_dim,
|
|
conv_alpha=conv_alpha,
|
|
block_dims=block_dims,
|
|
block_alphas=block_alphas,
|
|
conv_block_dims=conv_block_dims,
|
|
conv_block_alphas=conv_block_alphas,
|
|
varbose=True,
|
|
)
|
|
|
|
if up_lr_weight is not None or mid_lr_weight is not None or down_lr_weight is not None:
|
|
network.set_block_lr_weight(up_lr_weight, mid_lr_weight, down_lr_weight)
|
|
|
|
return network
|
|
|
|
|
|
# このメソッドは外部から呼び出される可能性を考慮しておく
|
|
# network_dim, network_alpha にはデフォルト値が入っている。
|
|
# block_dims, block_alphas は両方ともNoneまたは両方とも値が入っている
|
|
# conv_dim, conv_alpha は両方ともNoneまたは両方とも値が入っている
|
|
def get_block_dims_and_alphas(
|
|
block_dims, block_alphas, network_dim, network_alpha, conv_block_dims, conv_block_alphas, conv_dim, conv_alpha
|
|
):
|
|
num_total_blocks = LoRANetwork.NUM_OF_BLOCKS * 2 + 1
|
|
|
|
def parse_ints(s):
|
|
return [int(i) for i in s.split(",")]
|
|
|
|
def parse_floats(s):
|
|
return [float(i) for i in s.split(",")]
|
|
|
|
# block_dimsとblock_alphasをパースする。必ず値が入る
|
|
if block_dims is not None:
|
|
block_dims = parse_ints(block_dims)
|
|
assert (
|
|
len(block_dims) == num_total_blocks
|
|
), f"block_dims must have {num_total_blocks} elements / block_dimsは{num_total_blocks}個指定してください"
|
|
else:
|
|
print(f"block_dims is not specified. all dims are set to {network_dim} / block_dimsが指定されていません。すべてのdimは{network_dim}になります")
|
|
block_dims = [network_dim] * num_total_blocks
|
|
|
|
if block_alphas is not None:
|
|
block_alphas = parse_floats(block_alphas)
|
|
assert (
|
|
len(block_alphas) == num_total_blocks
|
|
), f"block_alphas must have {num_total_blocks} elements / block_alphasは{num_total_blocks}個指定してください"
|
|
else:
|
|
print(
|
|
f"block_alphas is not specified. all alphas are set to {network_alpha} / block_alphasが指定されていません。すべてのalphaは{network_alpha}になります"
|
|
)
|
|
block_alphas = [network_alpha] * num_total_blocks
|
|
|
|
# conv_block_dimsとconv_block_alphasを、指定がある場合のみパースする。指定がなければconv_dimとconv_alphaを使う
|
|
if conv_block_dims is not None:
|
|
conv_block_dims = parse_ints(conv_block_dims)
|
|
assert (
|
|
len(conv_block_dims) == num_total_blocks
|
|
), f"conv_block_dims must have {num_total_blocks} elements / conv_block_dimsは{num_total_blocks}個指定してください"
|
|
|
|
if conv_block_alphas is not None:
|
|
conv_block_alphas = parse_floats(conv_block_alphas)
|
|
assert (
|
|
len(conv_block_alphas) == num_total_blocks
|
|
), f"conv_block_alphas must have {num_total_blocks} elements / conv_block_alphasは{num_total_blocks}個指定してください"
|
|
else:
|
|
if conv_alpha is None:
|
|
conv_alpha = 1.0
|
|
print(
|
|
f"conv_block_alphas is not specified. all alphas are set to {conv_alpha} / conv_block_alphasが指定されていません。すべてのalphaは{conv_alpha}になります"
|
|
)
|
|
conv_block_alphas = [conv_alpha] * num_total_blocks
|
|
else:
|
|
if conv_dim is not None:
|
|
print(
|
|
f"conv_dim/alpha for all blocks are set to {conv_dim} and {conv_alpha} / すべてのブロックのconv_dimとalphaは{conv_dim}および{conv_alpha}になります"
|
|
)
|
|
conv_block_dims = [conv_dim] * num_total_blocks
|
|
conv_block_alphas = [conv_alpha] * num_total_blocks
|
|
else:
|
|
conv_block_dims = None
|
|
conv_block_alphas = None
|
|
|
|
return block_dims, block_alphas, conv_block_dims, conv_block_alphas
|
|
|
|
|
|
# 層別学習率用に層ごとの学習率に対する倍率を定義する、外部から呼び出される可能性を考慮しておく
|
|
def get_block_lr_weight(
|
|
down_lr_weight, mid_lr_weight, up_lr_weight, zero_threshold
|
|
) -> Tuple[List[float], List[float], List[float]]:
|
|
# パラメータ未指定時は何もせず、今までと同じ動作とする
|
|
if up_lr_weight is None and mid_lr_weight is None and down_lr_weight is None:
|
|
return None, None, None
|
|
|
|
max_len = LoRANetwork.NUM_OF_BLOCKS # フルモデル相当でのup,downの層の数
|
|
|
|
def get_list(name_with_suffix) -> List[float]:
|
|
import math
|
|
|
|
tokens = name_with_suffix.split("+")
|
|
name = tokens[0]
|
|
base_lr = float(tokens[1]) if len(tokens) > 1 else 0.0
|
|
|
|
if name == "cosine":
|
|
return [math.sin(math.pi * (i / (max_len - 1)) / 2) + base_lr for i in reversed(range(max_len))]
|
|
elif name == "sine":
|
|
return [math.sin(math.pi * (i / (max_len - 1)) / 2) + base_lr for i in range(max_len)]
|
|
elif name == "linear":
|
|
return [i / (max_len - 1) + base_lr for i in range(max_len)]
|
|
elif name == "reverse_linear":
|
|
return [i / (max_len - 1) + base_lr for i in reversed(range(max_len))]
|
|
elif name == "zeros":
|
|
return [0.0 + base_lr] * max_len
|
|
else:
|
|
print(
|
|
"Unknown lr_weight argument %s is used. Valid arguments: / 不明なlr_weightの引数 %s が使われました。有効な引数:\n\tcosine, sine, linear, reverse_linear, zeros"
|
|
% (name)
|
|
)
|
|
return None
|
|
|
|
if type(down_lr_weight) == str:
|
|
down_lr_weight = get_list(down_lr_weight)
|
|
if type(up_lr_weight) == str:
|
|
up_lr_weight = get_list(up_lr_weight)
|
|
|
|
if (up_lr_weight != None and len(up_lr_weight) > max_len) or (down_lr_weight != None and len(down_lr_weight) > max_len):
|
|
print("down_weight or up_weight is too long. Parameters after %d-th are ignored." % max_len)
|
|
print("down_weightもしくはup_weightが長すぎます。%d個目以降のパラメータは無視されます。" % max_len)
|
|
up_lr_weight = up_lr_weight[:max_len]
|
|
down_lr_weight = down_lr_weight[:max_len]
|
|
|
|
if (up_lr_weight != None and len(up_lr_weight) < max_len) or (down_lr_weight != None and len(down_lr_weight) < max_len):
|
|
print("down_weight or up_weight is too short. Parameters after %d-th are filled with 1." % max_len)
|
|
print("down_weightもしくはup_weightが短すぎます。%d個目までの不足したパラメータは1で補われます。" % max_len)
|
|
|
|
if down_lr_weight != None and len(down_lr_weight) < max_len:
|
|
down_lr_weight = down_lr_weight + [1.0] * (max_len - len(down_lr_weight))
|
|
if up_lr_weight != None and len(up_lr_weight) < max_len:
|
|
up_lr_weight = up_lr_weight + [1.0] * (max_len - len(up_lr_weight))
|
|
|
|
if (up_lr_weight != None) or (mid_lr_weight != None) or (down_lr_weight != None):
|
|
print("apply block learning rate / 階層別学習率を適用します。")
|
|
if down_lr_weight != None:
|
|
down_lr_weight = [w if w > zero_threshold else 0 for w in down_lr_weight]
|
|
print("down_lr_weight (shallower -> deeper, 浅い層->深い層):", down_lr_weight)
|
|
else:
|
|
print("down_lr_weight: all 1.0, すべて1.0")
|
|
|
|
if mid_lr_weight != None:
|
|
mid_lr_weight = mid_lr_weight if mid_lr_weight > zero_threshold else 0
|
|
print("mid_lr_weight:", mid_lr_weight)
|
|
else:
|
|
print("mid_lr_weight: 1.0")
|
|
|
|
if up_lr_weight != None:
|
|
up_lr_weight = [w if w > zero_threshold else 0 for w in up_lr_weight]
|
|
print("up_lr_weight (deeper -> shallower, 深い層->浅い層):", up_lr_weight)
|
|
else:
|
|
print("up_lr_weight: all 1.0, すべて1.0")
|
|
|
|
return down_lr_weight, mid_lr_weight, up_lr_weight
|
|
|
|
|
|
# lr_weightが0のblockをblock_dimsから除外する、外部から呼び出す可能性を考慮しておく
|
|
def remove_block_dims_and_alphas(
|
|
block_dims, block_alphas, conv_block_dims, conv_block_alphas, down_lr_weight, mid_lr_weight, up_lr_weight
|
|
):
|
|
# set 0 to block dim without learning rate to remove the block
|
|
if down_lr_weight != None:
|
|
for i, lr in enumerate(down_lr_weight):
|
|
if lr == 0:
|
|
block_dims[i] = 0
|
|
if conv_block_dims is not None:
|
|
conv_block_dims[i] = 0
|
|
if mid_lr_weight != None:
|
|
if mid_lr_weight == 0:
|
|
block_dims[LoRANetwork.NUM_OF_BLOCKS] = 0
|
|
if conv_block_dims is not None:
|
|
conv_block_dims[LoRANetwork.NUM_OF_BLOCKS] = 0
|
|
if up_lr_weight != None:
|
|
for i, lr in enumerate(up_lr_weight):
|
|
if lr == 0:
|
|
block_dims[LoRANetwork.NUM_OF_BLOCKS + 1 + i] = 0
|
|
if conv_block_dims is not None:
|
|
conv_block_dims[LoRANetwork.NUM_OF_BLOCKS + 1 + i] = 0
|
|
|
|
return block_dims, block_alphas, conv_block_dims, conv_block_alphas
|
|
|
|
|
|
# 外部から呼び出す可能性を考慮しておく
|
|
def get_block_index(lora_name: str) -> int:
|
|
block_idx = -1 # invalid lora name
|
|
|
|
m = RE_UPDOWN.search(lora_name)
|
|
if m:
|
|
g = m.groups()
|
|
i = int(g[1])
|
|
j = int(g[3])
|
|
if g[2] == "resnets":
|
|
idx = 3 * i + j
|
|
elif g[2] == "attentions":
|
|
idx = 3 * i + j
|
|
elif g[2] == "upsamplers" or g[2] == "downsamplers":
|
|
idx = 3 * i + 2
|
|
|
|
if g[0] == "down":
|
|
block_idx = 1 + idx # 0に該当するLoRAは存在しない
|
|
elif g[0] == "up":
|
|
block_idx = LoRANetwork.NUM_OF_BLOCKS + 1 + idx
|
|
|
|
elif "mid_block_" in lora_name:
|
|
block_idx = LoRANetwork.NUM_OF_BLOCKS # idx=12
|
|
|
|
return block_idx
|
|
|
|
|
|
# Create network from weights for inference, weights are not loaded here (because can be merged)
|
|
def create_network_from_weights(multiplier, file, vae, text_encoder, unet, weights_sd=None, **kwargs):
|
|
if weights_sd is None:
|
|
if os.path.splitext(file)[1] == ".safetensors":
|
|
from safetensors.torch import load_file, safe_open
|
|
|
|
weights_sd = load_file(file)
|
|
else:
|
|
weights_sd = torch.load(file, map_location="cpu")
|
|
|
|
# get dim/alpha mapping
|
|
modules_dim = {}
|
|
modules_alpha = {}
|
|
for key, value in weights_sd.items():
|
|
if "." not in key:
|
|
continue
|
|
|
|
lora_name = key.split(".")[0]
|
|
if "alpha" in key:
|
|
modules_alpha[lora_name] = value
|
|
elif "lora_down" in key:
|
|
dim = value.size()[0]
|
|
modules_dim[lora_name] = dim
|
|
# print(lora_name, value.size(), dim)
|
|
|
|
# support old LoRA without alpha
|
|
for key in modules_dim.keys():
|
|
if key not in modules_alpha:
|
|
modules_alpha = modules_dim[key]
|
|
|
|
network = LoRANetwork(text_encoder, unet, multiplier=multiplier, modules_dim=modules_dim, modules_alpha=modules_alpha)
|
|
return network, weights_sd
|
|
|
|
|
|
class LoRANetwork(torch.nn.Module):
|
|
NUM_OF_BLOCKS = 12 # フルモデル相当でのup,downの層の数
|
|
|
|
# is it possible to apply conv_in and conv_out? -> yes, newer LoCon supports it (^^;)
|
|
UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"]
|
|
UNET_TARGET_REPLACE_MODULE_CONV2D_3X3 = ["ResnetBlock2D", "Downsample2D", "Upsample2D"]
|
|
TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
|
|
LORA_PREFIX_UNET = "lora_unet"
|
|
LORA_PREFIX_TEXT_ENCODER = "lora_te"
|
|
|
|
def __init__(
|
|
self,
|
|
text_encoder,
|
|
unet,
|
|
multiplier=1.0,
|
|
lora_dim=4,
|
|
alpha=1,
|
|
conv_lora_dim=None,
|
|
conv_alpha=None,
|
|
block_dims=None,
|
|
block_alphas=None,
|
|
conv_block_dims=None,
|
|
conv_block_alphas=None,
|
|
modules_dim=None,
|
|
modules_alpha=None,
|
|
varbose=False,
|
|
) -> None:
|
|
"""
|
|
LoRA network: すごく引数が多いが、パターンは以下の通り
|
|
1. lora_dimとalphaを指定
|
|
2. lora_dim、alpha、conv_lora_dim、conv_alphaを指定
|
|
3. block_dimsとblock_alphasを指定 : Conv2d3x3には適用しない
|
|
4. block_dims、block_alphas、conv_block_dims、conv_block_alphasを指定 : Conv2d3x3にも適用する
|
|
5. modules_dimとmodules_alphaを指定 (推論用)
|
|
"""
|
|
super().__init__()
|
|
self.multiplier = multiplier
|
|
|
|
self.lora_dim = lora_dim
|
|
self.alpha = alpha
|
|
self.conv_lora_dim = conv_lora_dim
|
|
self.conv_alpha = conv_alpha
|
|
|
|
if modules_dim is not None:
|
|
print(f"create LoRA network from weights")
|
|
elif block_dims is not None:
|
|
print(f"create LoRA network from block_dims")
|
|
print(f"block_dims: {block_dims}")
|
|
print(f"block_alphas: {block_alphas}")
|
|
if conv_block_dims is not None:
|
|
print(f"conv_block_dims: {conv_block_dims}")
|
|
print(f"conv_block_alphas: {conv_block_alphas}")
|
|
else:
|
|
print(f"create LoRA network. base dim (rank): {lora_dim}, alpha: {alpha}")
|
|
if self.conv_lora_dim is not None:
|
|
print(f"apply LoRA to Conv2d with kernel size (3,3). dim (rank): {self.conv_lora_dim}, alpha: {self.conv_alpha}")
|
|
|
|
# create module instances
|
|
def create_modules(is_unet, root_module: torch.nn.Module, target_replace_modules) -> List[LoRAModule]:
|
|
prefix = LoRANetwork.LORA_PREFIX_UNET if is_unet else LoRANetwork.LORA_PREFIX_TEXT_ENCODER
|
|
loras = []
|
|
skipped = []
|
|
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():
|
|
is_linear = child_module.__class__.__name__ == "Linear"
|
|
is_conv2d = child_module.__class__.__name__ == "Conv2d"
|
|
is_conv2d_1x1 = is_conv2d and child_module.kernel_size == (1, 1)
|
|
|
|
if is_linear or is_conv2d:
|
|
lora_name = prefix + "." + name + "." + child_name
|
|
lora_name = lora_name.replace(".", "_")
|
|
|
|
dim = None
|
|
alpha = None
|
|
if modules_dim is not None:
|
|
if lora_name in modules_dim:
|
|
dim = modules_dim[lora_name]
|
|
alpha = modules_alpha[lora_name]
|
|
elif is_unet and block_dims is not None:
|
|
block_idx = get_block_index(lora_name)
|
|
if is_linear or is_conv2d_1x1:
|
|
dim = block_dims[block_idx]
|
|
alpha = block_alphas[block_idx]
|
|
elif conv_block_dims is not None:
|
|
dim = conv_block_dims[block_idx]
|
|
alpha = conv_block_alphas[block_idx]
|
|
else:
|
|
if is_linear or is_conv2d_1x1:
|
|
dim = self.lora_dim
|
|
alpha = self.alpha
|
|
elif self.conv_lora_dim is not None:
|
|
dim = self.conv_lora_dim
|
|
alpha = self.conv_alpha
|
|
|
|
if dim is None or dim == 0:
|
|
if is_linear or is_conv2d_1x1 or (self.conv_lora_dim is not None or conv_block_dims is not None):
|
|
skipped.append(lora_name)
|
|
continue
|
|
|
|
lora = LoRAModule(lora_name, child_module, self.multiplier, dim, alpha)
|
|
loras.append(lora)
|
|
return loras, skipped
|
|
|
|
self.text_encoder_loras, skipped_te = create_modules(False, text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE)
|
|
print(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.")
|
|
|
|
# extend U-Net target modules if conv2d 3x3 is enabled, or load from weights
|
|
target_modules = LoRANetwork.UNET_TARGET_REPLACE_MODULE
|
|
if modules_dim is not None or self.conv_lora_dim is not None or conv_block_dims is not None:
|
|
target_modules += LoRANetwork.UNET_TARGET_REPLACE_MODULE_CONV2D_3X3
|
|
|
|
self.unet_loras, skipped_un = create_modules(True, unet, target_modules)
|
|
print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
|
|
|
|
skipped = skipped_te + skipped_un
|
|
if varbose and len(skipped) > 0:
|
|
print(
|
|
f"because block_lr_weight is 0 or dim (rank) is 0, {len(skipped)} LoRA modules are skipped / block_lr_weightまたはdim (rank)が0の為、次の{len(skipped)}個のLoRAモジュールはスキップされます:"
|
|
)
|
|
for name in skipped:
|
|
print(f"\t{name}")
|
|
|
|
self.up_lr_weight: List[float] = None
|
|
self.down_lr_weight: List[float] = None
|
|
self.mid_lr_weight: float = None
|
|
self.block_lr = False
|
|
|
|
# assertion
|
|
names = set()
|
|
for lora in self.text_encoder_loras + self.unet_loras:
|
|
assert lora.lora_name not in names, f"duplicated lora name: {lora.lora_name}"
|
|
names.add(lora.lora_name)
|
|
|
|
def set_multiplier(self, multiplier):
|
|
self.multiplier = multiplier
|
|
for lora in self.text_encoder_loras + self.unet_loras:
|
|
lora.multiplier = self.multiplier
|
|
|
|
def load_weights(self, file):
|
|
if os.path.splitext(file)[1] == ".safetensors":
|
|
from safetensors.torch import load_file
|
|
|
|
weights_sd = load_file(file)
|
|
else:
|
|
weights_sd = torch.load(file, map_location="cpu")
|
|
|
|
info = self.load_state_dict(weights_sd, False)
|
|
return info
|
|
|
|
def apply_to(self, text_encoder, unet, apply_text_encoder=True, apply_unet=True):
|
|
if apply_text_encoder:
|
|
print("enable LoRA for text encoder")
|
|
else:
|
|
self.text_encoder_loras = []
|
|
|
|
if apply_unet:
|
|
print("enable LoRA for U-Net")
|
|
else:
|
|
self.unet_loras = []
|
|
|
|
for lora in self.text_encoder_loras + self.unet_loras:
|
|
lora.apply_to()
|
|
self.add_module(lora.lora_name, lora)
|
|
|
|
# TODO refactor to common function with apply_to
|
|
def merge_to(self, text_encoder, unet, weights_sd, dtype, device):
|
|
apply_text_encoder = apply_unet = False
|
|
for key in weights_sd.keys():
|
|
if key.startswith(LoRANetwork.LORA_PREFIX_TEXT_ENCODER):
|
|
apply_text_encoder = True
|
|
elif key.startswith(LoRANetwork.LORA_PREFIX_UNET):
|
|
apply_unet = True
|
|
|
|
if apply_text_encoder:
|
|
print("enable LoRA for text encoder")
|
|
else:
|
|
self.text_encoder_loras = []
|
|
|
|
if apply_unet:
|
|
print("enable LoRA for U-Net")
|
|
else:
|
|
self.unet_loras = []
|
|
|
|
for lora in self.text_encoder_loras + self.unet_loras:
|
|
sd_for_lora = {}
|
|
for key in weights_sd.keys():
|
|
if key.startswith(lora.lora_name):
|
|
sd_for_lora[key[len(lora.lora_name) + 1 :]] = weights_sd[key]
|
|
lora.merge_to(sd_for_lora, dtype, device)
|
|
|
|
print(f"weights are merged")
|
|
|
|
# 層別学習率用に層ごとの学習率に対する倍率を定義する
|
|
def set_block_lr_weight(
|
|
self,
|
|
up_lr_weight: List[float] = None,
|
|
mid_lr_weight: float = None,
|
|
down_lr_weight: List[float] = None,
|
|
):
|
|
self.block_lr = True
|
|
self.down_lr_weight = down_lr_weight
|
|
self.mid_lr_weight = mid_lr_weight
|
|
self.up_lr_weight = up_lr_weight
|
|
|
|
def get_lr_weight(self, lora: LoRAModule) -> float:
|
|
lr_weight = 1.0
|
|
block_idx = get_block_index(lora.lora_name)
|
|
if block_idx < 0:
|
|
return lr_weight
|
|
|
|
if block_idx < LoRANetwork.NUM_OF_BLOCKS:
|
|
if self.down_lr_weight != None:
|
|
lr_weight = self.down_lr_weight[block_idx]
|
|
elif block_idx == LoRANetwork.NUM_OF_BLOCKS:
|
|
if self.mid_lr_weight != None:
|
|
lr_weight = self.mid_lr_weight
|
|
elif block_idx > LoRANetwork.NUM_OF_BLOCKS:
|
|
if self.up_lr_weight != None:
|
|
lr_weight = self.up_lr_weight[block_idx - LoRANetwork.NUM_OF_BLOCKS - 1]
|
|
|
|
return lr_weight
|
|
|
|
def prepare_optimizer_params(self, text_encoder_lr, unet_lr, default_lr):
|
|
self.requires_grad_(True)
|
|
all_params = []
|
|
|
|
def enumerate_params(loras):
|
|
params = []
|
|
for lora in loras:
|
|
params.extend(lora.parameters())
|
|
return params
|
|
|
|
if self.text_encoder_loras:
|
|
param_data = {"params": enumerate_params(self.text_encoder_loras)}
|
|
if text_encoder_lr is not None:
|
|
param_data["lr"] = text_encoder_lr
|
|
all_params.append(param_data)
|
|
|
|
if self.unet_loras:
|
|
if self.block_lr:
|
|
# 学習率のグラフをblockごとにしたいので、blockごとにloraを分類
|
|
block_idx_to_lora = {}
|
|
for lora in self.unet_loras:
|
|
idx = get_block_index(lora.lora_name)
|
|
if idx not in block_idx_to_lora:
|
|
block_idx_to_lora[idx] = []
|
|
block_idx_to_lora[idx].append(lora)
|
|
|
|
# blockごとにパラメータを設定する
|
|
for idx, block_loras in block_idx_to_lora.items():
|
|
param_data = {"params": enumerate_params(block_loras)}
|
|
|
|
if unet_lr is not None:
|
|
param_data["lr"] = unet_lr * self.get_lr_weight(block_loras[0])
|
|
elif default_lr is not None:
|
|
param_data["lr"] = default_lr * self.get_lr_weight(block_loras[0])
|
|
if ("lr" in param_data) and (param_data["lr"] == 0):
|
|
continue
|
|
all_params.append(param_data)
|
|
|
|
else:
|
|
param_data = {"params": enumerate_params(self.unet_loras)}
|
|
if unet_lr is not None:
|
|
param_data["lr"] = unet_lr
|
|
all_params.append(param_data)
|
|
|
|
return all_params
|
|
|
|
def enable_gradient_checkpointing(self):
|
|
# not supported
|
|
pass
|
|
|
|
def prepare_grad_etc(self, text_encoder, unet):
|
|
self.requires_grad_(True)
|
|
|
|
def on_epoch_start(self, text_encoder, unet):
|
|
self.train()
|
|
|
|
def get_trainable_params(self):
|
|
return self.parameters()
|
|
|
|
def save_weights(self, file, dtype, metadata):
|
|
if metadata is not None and len(metadata) == 0:
|
|
metadata = None
|
|
|
|
state_dict = self.state_dict()
|
|
|
|
if dtype is not None:
|
|
for key in list(state_dict.keys()):
|
|
v = state_dict[key]
|
|
v = v.detach().clone().to("cpu").to(dtype)
|
|
state_dict[key] = v
|
|
|
|
if os.path.splitext(file)[1] == ".safetensors":
|
|
from safetensors.torch import save_file
|
|
|
|
# Precalculate model hashes to save time on indexing
|
|
if metadata is None:
|
|
metadata = {}
|
|
model_hash, legacy_hash = train_util.precalculate_safetensors_hashes(state_dict, metadata)
|
|
metadata["sshs_model_hash"] = model_hash
|
|
metadata["sshs_legacy_hash"] = legacy_hash
|
|
|
|
save_file(state_dict, file, metadata)
|
|
else:
|
|
torch.save(state_dict, file)
|
|
|
|
@staticmethod
|
|
def set_regions(networks, image):
|
|
image = image.astype(np.float32) / 255.0
|
|
for i, network in enumerate(networks[:3]):
|
|
# NOTE: consider averaging overwrapping area
|
|
region = image[:, :, i]
|
|
if region.max() == 0:
|
|
continue
|
|
region = torch.tensor(region)
|
|
network.set_region(region)
|
|
|
|
def set_region(self, region):
|
|
for lora in self.unet_loras:
|
|
lora.set_region(region)
|