VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-08 22:35:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

support block dim/lr for sdxl

Browse Source
This commit is contained in:
Kohya S
2024-05-03 22:18:20 +09:00
parent 969f82ab47
commit 58c2d856ae
2 changed files with 156 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files

237
networks/lora.py
View File

@@ -12,6 +12,7 @@ import numpy as np
import torch import torch
import re import re
from library.utils import setup_logging from library.utils import setup_logging
from library.sdxl_original_unet import SdxlUNet2DConditionModel
setup_logging() setup_logging()
import logging import logging
@@ -385,14 +386,14 @@ class LoRAInfModule(LoRAModule):
return out return out
def parse_block_lr_kwargs(nw_kwargs): def parse_block_lr_kwargs(is_sdxl: bool, nw_kwargs: Dict) -> Optional[List[float]]:
down_lr_weight = nw_kwargs.get("down_lr_weight", None) down_lr_weight = nw_kwargs.get("down_lr_weight", None)
mid_lr_weight = nw_kwargs.get("mid_lr_weight", None) mid_lr_weight = nw_kwargs.get("mid_lr_weight", None)
up_lr_weight = nw_kwargs.get("up_lr_weight", None) up_lr_weight = nw_kwargs.get("up_lr_weight", None)
# 以上のいずれにも設定がない場合は無効としてNoneを返す # 以上のいずれにも設定がない場合は無効としてNoneを返す
if down_lr_weight is None and mid_lr_weight is None and up_lr_weight is None: if down_lr_weight is None and mid_lr_weight is None and up_lr_weight is None:
return None, None, None return None
# extract learning rate weight for each block # extract learning rate weight for each block
if down_lr_weight is not None: if down_lr_weight is not None:
@@ -401,18 +402,16 @@ def parse_block_lr_kwargs(nw_kwargs):
down_lr_weight = [(float(s) if s else 0.0) for s in down_lr_weight.split(",")] down_lr_weight = [(float(s) if s else 0.0) for s in down_lr_weight.split(",")]
if mid_lr_weight is not None: if mid_lr_weight is not None:
mid_lr_weight = float(mid_lr_weight) mid_lr_weight = [(float(s) if s else 0.0) for s in mid_lr_weight.split(",")]
if up_lr_weight is not None: if up_lr_weight is not None:
if "," in up_lr_weight: if "," in up_lr_weight:
up_lr_weight = [(float(s) if s else 0.0) for s in up_lr_weight.split(",")] 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( return get_block_lr_weight(
down_lr_weight, mid_lr_weight, up_lr_weight, float(nw_kwargs.get("block_lr_zero_threshold", 0.0)) is_sdxl, down_lr_weight, mid_lr_weight, up_lr_weight, float(nw_kwargs.get("block_lr_zero_threshold", 0.0))
) )
return down_lr_weight, mid_lr_weight, up_lr_weight
def create_network( def create_network(
multiplier: float, multiplier: float,
@@ -424,6 +423,9 @@ def create_network(
neuron_dropout: Optional[float] = None, neuron_dropout: Optional[float] = None,
**kwargs, **kwargs,
): ):
# if unet is an instance of SdxlUNet2DConditionModel or subclass, set is_sdxl to True
is_sdxl = unet is not None and issubclass(unet.__class__, SdxlUNet2DConditionModel)
if network_dim is None: if network_dim is None:
network_dim = 4 # default network_dim = 4 # default
if network_alpha is None: if network_alpha is None:
@@ -441,21 +443,21 @@ def create_network(
# block dim/alpha/lr # block dim/alpha/lr
block_dims = kwargs.get("block_dims", None) block_dims = kwargs.get("block_dims", None)
down_lr_weight, mid_lr_weight, up_lr_weight = parse_block_lr_kwargs(kwargs) block_lr_weight = parse_block_lr_kwargs(is_sdxl, kwargs)
# 以上のいずれかに指定があればblockごとのdim(rank)を有効にする # 以上のいずれかに指定があれば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: if block_dims is not None or block_lr_weight is not None:
block_alphas = kwargs.get("block_alphas", None) block_alphas = kwargs.get("block_alphas", None)
conv_block_dims = kwargs.get("conv_block_dims", None) conv_block_dims = kwargs.get("conv_block_dims", None)
conv_block_alphas = kwargs.get("conv_block_alphas", 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, 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 is_sdxl, block_dims, block_alphas, network_dim, network_alpha, conv_block_dims, conv_block_alphas, conv_dim, conv_alpha
) )
# remove block dim/alpha without learning rate # 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 = remove_block_dims_and_alphas(
block_dims, block_alphas, conv_block_dims, conv_block_alphas, down_lr_weight, mid_lr_weight, up_lr_weight is_sdxl, block_dims, block_alphas, conv_block_dims, conv_block_alphas, block_lr_weight
) )
else: else:
@@ -488,6 +490,7 @@ def create_network(
conv_block_dims=conv_block_dims, conv_block_dims=conv_block_dims,
conv_block_alphas=conv_block_alphas, conv_block_alphas=conv_block_alphas,
varbose=True, varbose=True,
is_sdxl=is_sdxl,
) )
loraplus_lr_ratio = kwargs.get("loraplus_lr_ratio", None) loraplus_lr_ratio = kwargs.get("loraplus_lr_ratio", None)
@@ -498,8 +501,8 @@ def create_network(
loraplus_text_encoder_lr_ratio = float(loraplus_text_encoder_lr_ratio) if loraplus_text_encoder_lr_ratio is not None else None loraplus_text_encoder_lr_ratio = float(loraplus_text_encoder_lr_ratio) if loraplus_text_encoder_lr_ratio is not None else None
network.set_loraplus_lr_ratio(loraplus_lr_ratio, loraplus_unet_lr_ratio, loraplus_text_encoder_lr_ratio) network.set_loraplus_lr_ratio(loraplus_lr_ratio, loraplus_unet_lr_ratio, loraplus_text_encoder_lr_ratio)
if up_lr_weight is not None or mid_lr_weight is not None or down_lr_weight is not None: if block_lr_weight is not None:
network.set_block_lr_weight(up_lr_weight, mid_lr_weight, down_lr_weight) network.set_block_lr_weight(block_lr_weight)
return network return network
@@ -509,9 +512,13 @@ def create_network(
# block_dims, block_alphas は両方ともNoneまたは両方とも値が入っている # block_dims, block_alphas は両方ともNoneまたは両方とも値が入っている
# conv_dim, conv_alpha は両方ともNoneまたは両方とも値が入っている # conv_dim, conv_alpha は両方ともNoneまたは両方とも値が入っている
def get_block_dims_and_alphas( def get_block_dims_and_alphas(
block_dims, block_alphas, network_dim, network_alpha, conv_block_dims, conv_block_alphas, conv_dim, conv_alpha is_sdxl, 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 if not is_sdxl:
num_total_blocks = LoRANetwork.NUM_OF_BLOCKS * 2 + LoRANetwork.NUM_OF_MID_BLOCKS
else:
# 1+9+3+9+1=23, no LoRA for emb_layers (0)
num_total_blocks = 1 + LoRANetwork.SDXL_NUM_OF_BLOCKS * 2 + LoRANetwork.SDXL_NUM_OF_MID_BLOCKS + 1
def parse_ints(s): def parse_ints(s):
return [int(i) for i in s.split(",")] return [int(i) for i in s.split(",")]
@@ -522,9 +529,10 @@ def get_block_dims_and_alphas(
# block_dimsとblock_alphasをパースする。必ず値が入る # block_dimsとblock_alphasをパースする。必ず値が入る
if block_dims is not None: if block_dims is not None:
block_dims = parse_ints(block_dims) block_dims = parse_ints(block_dims)
assert ( assert len(block_dims) == num_total_blocks, (
len(block_dims) == num_total_blocks f"block_dims must have {num_total_blocks} elements but {len(block_dims)} elements are given"
), f"block_dims must have {num_total_blocks} elements / block_dimsは{num_total_blocks}個指定してください" + f" / block_dimsは{num_total_blocks}個指定してください（指定された個数: {len(block_dims)}）"
)
else: else:
logger.warning( logger.warning(
f"block_dims is not specified. all dims are set to {network_dim} / block_dimsが指定されていません。すべてのdimは{network_dim}になります" f"block_dims is not specified. all dims are set to {network_dim} / block_dimsが指定されていません。すべてのdimは{network_dim}になります"
@@ -575,15 +583,25 @@ def get_block_dims_and_alphas(
return block_dims, block_alphas, conv_block_dims, conv_block_alphas return block_dims, block_alphas, conv_block_dims, conv_block_alphas
# 層別学習率用に層ごとの学習率に対する倍率を定義する、外部から呼び出される可能性を考慮しておく # 層別学習率用に層ごとの学習率に対する倍率を定義する、外部から呼び出せるようにclass外に出しておく
# 戻り値は block ごとの倍率のリスト
def get_block_lr_weight( def get_block_lr_weight(
down_lr_weight, mid_lr_weight, up_lr_weight, zero_threshold is_sdxl,
) -> Tuple[List[float], List[float], List[float]]: down_lr_weight: Union[str, List[float]],
mid_lr_weight: List[float],
up_lr_weight: Union[str, List[float]],
zero_threshold: float,
) -> Optional[List[float]]:
# パラメータ未指定時は何もせず、今までと同じ動作とする # パラメータ未指定時は何もせず、今までと同じ動作とする
if up_lr_weight is None and mid_lr_weight is None and down_lr_weight is None: if up_lr_weight is None and mid_lr_weight is None and down_lr_weight is None:
return None, None, None return None
max_len = LoRANetwork.NUM_OF_BLOCKS # フルモデル相当でのup,downの層の数 if not is_sdxl:
max_len_for_down_or_up = LoRANetwork.NUM_OF_BLOCKS
max_len_for_mid = LoRANetwork.NUM_OF_MID_BLOCKS
else:
max_len_for_down_or_up = LoRANetwork.SDXL_NUM_OF_BLOCKS
max_len_for_mid = LoRANetwork.SDXL_NUM_OF_MID_BLOCKS
def get_list(name_with_suffix) -> List[float]: def get_list(name_with_suffix) -> List[float]:
import math import math
@@ -593,15 +611,18 @@ def get_block_lr_weight(
base_lr = float(tokens[1]) if len(tokens) > 1 else 0.0 base_lr = float(tokens[1]) if len(tokens) > 1 else 0.0
if name == "cosine": if name == "cosine":
return [math.sin(math.pi * (i / (max_len - 1)) / 2) + base_lr for i in reversed(range(max_len))] return [
math.sin(math.pi * (i / (max_len_for_down_or_up - 1)) / 2) + base_lr
for i in reversed(range(max_len_for_down_or_up))
]
elif name == "sine": elif name == "sine":
return [math.sin(math.pi * (i / (max_len - 1)) / 2) + base_lr for i in range(max_len)] return [math.sin(math.pi * (i / (max_len_for_down_or_up - 1)) / 2) + base_lr for i in range(max_len_for_down_or_up)]
elif name == "linear": elif name == "linear":
return [i / (max_len - 1) + base_lr for i in range(max_len)] return [i / (max_len_for_down_or_up - 1) + base_lr for i in range(max_len_for_down_or_up)]
elif name == "reverse_linear": elif name == "reverse_linear":
return [i / (max_len - 1) + base_lr for i in reversed(range(max_len))] return [i / (max_len_for_down_or_up - 1) + base_lr for i in reversed(range(max_len_for_down_or_up))]
elif name == "zeros": elif name == "zeros":
return [0.0 + base_lr] * max_len return [0.0 + base_lr] * max_len_for_down_or_up
else: else:
logger.error( logger.error(
"Unknown lr_weight argument %s is used. Valid arguments: / 不明なlr_weightの引数 %s が使われました。有効な引数:\n\tcosine, sine, linear, reverse_linear, zeros" "Unknown lr_weight argument %s is used. Valid arguments: / 不明なlr_weightの引数 %s が使われました。有効な引数:\n\tcosine, sine, linear, reverse_linear, zeros"
@@ -614,20 +635,36 @@ def get_block_lr_weight(
if type(up_lr_weight) == str: if type(up_lr_weight) == str:
up_lr_weight = get_list(up_lr_weight) 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): if (up_lr_weight != None and len(up_lr_weight) > max_len_for_down_or_up) or (
logger.warning("down_weight or up_weight is too long. Parameters after %d-th are ignored." % max_len) down_lr_weight != None and len(down_lr_weight) > max_len_for_down_or_up
logger.warning("down_weightもしくはup_weightが長すぎます。%d個目以降のパラメータは無視されます。" % max_len) ):
up_lr_weight = up_lr_weight[:max_len] logger.warning("down_weight or up_weight is too long. Parameters after %d-th are ignored." % max_len_for_down_or_up)
down_lr_weight = down_lr_weight[:max_len] logger.warning("down_weightもしくはup_weightが長すぎます。%d個目以降のパラメータは無視されます。" % max_len_for_down_or_up)
up_lr_weight = up_lr_weight[:max_len_for_down_or_up]
down_lr_weight = down_lr_weight[:max_len_for_down_or_up]
if (up_lr_weight != None and len(up_lr_weight) < max_len) or (down_lr_weight != None and len(down_lr_weight) < max_len): if mid_lr_weight != None and len(mid_lr_weight) > max_len_for_mid:
logger.warning("down_weight or up_weight is too short. Parameters after %d-th are filled with 1." % max_len) logger.warning("mid_weight is too long. Parameters after %d-th are ignored." % max_len_for_mid)
logger.warning("down_weightもしくはup_weightが短すぎます。%d個目までの不足したパラメータは1で補われます。" % max_len) logger.warning("mid_weightが長すぎます。%d個目以降のパラメータは無視されます。" % max_len_for_mid)
mid_lr_weight = mid_lr_weight[:max_len_for_mid]
if down_lr_weight != None and len(down_lr_weight) < max_len: if (up_lr_weight != None and len(up_lr_weight) < max_len_for_down_or_up) or (
down_lr_weight = down_lr_weight + [1.0] * (max_len - len(down_lr_weight)) down_lr_weight != None and len(down_lr_weight) < max_len_for_down_or_up
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)) logger.warning("down_weight or up_weight is too short. Parameters after %d-th are filled with 1." % max_len_for_down_or_up)
logger.warning(
"down_weightもしくはup_weightが短すぎます。%d個目までの不足したパラメータは1で補われます。" % max_len_for_down_or_up
)
if down_lr_weight != None and len(down_lr_weight) < max_len_for_down_or_up:
down_lr_weight = down_lr_weight + [1.0] * (max_len_for_down_or_up - len(down_lr_weight))
if up_lr_weight != None and len(up_lr_weight) < max_len_for_down_or_up:
up_lr_weight = up_lr_weight + [1.0] * (max_len_for_down_or_up - len(up_lr_weight))
if mid_lr_weight != None and len(mid_lr_weight) < max_len_for_mid:
logger.warning("mid_weight is too short. Parameters after %d-th are filled with 1." % max_len_for_mid)
logger.warning("mid_weightが短すぎます。%d個目までの不足したパラメータは1で補われます。" % max_len_for_mid)
mid_lr_weight = mid_lr_weight + [1.0] * (max_len_for_mid - len(mid_lr_weight))
if (up_lr_weight != None) or (mid_lr_weight != None) or (down_lr_weight != None): if (up_lr_weight != None) or (mid_lr_weight != None) or (down_lr_weight != None):
logger.info("apply block learning rate / 階層別学習率を適用します。") logger.info("apply block learning rate / 階層別学習率を適用します。")
@@ -635,53 +672,52 @@ def get_block_lr_weight(
down_lr_weight = [w if w > zero_threshold else 0 for w in down_lr_weight] down_lr_weight = [w if w > zero_threshold else 0 for w in down_lr_weight]
logger.info(f"down_lr_weight (shallower -> deeper, 浅い層->深い層): {down_lr_weight}") logger.info(f"down_lr_weight (shallower -> deeper, 浅い層->深い層): {down_lr_weight}")
else: else:
down_lr_weight = [1.0] * max_len_for_down_or_up
logger.info("down_lr_weight: all 1.0, すべて1.0") logger.info("down_lr_weight: all 1.0, すべて1.0")
if mid_lr_weight != None: if mid_lr_weight != None:
mid_lr_weight = mid_lr_weight if mid_lr_weight > zero_threshold else 0 mid_lr_weight = [w if w > zero_threshold else 0 for w in mid_lr_weight]
logger.info(f"mid_lr_weight: {mid_lr_weight}") logger.info(f"mid_lr_weight: {mid_lr_weight}")
else: else:
logger.info("mid_lr_weight: 1.0") mid_lr_weight = [1.0] * max_len_for_mid
logger.info("mid_lr_weight: all 1.0, すべて1.0")
if up_lr_weight != None: if up_lr_weight != None:
up_lr_weight = [w if w > zero_threshold else 0 for w in up_lr_weight] up_lr_weight = [w if w > zero_threshold else 0 for w in up_lr_weight]
logger.info(f"up_lr_weight (deeper -> shallower, 深い層->浅い層): {up_lr_weight}") logger.info(f"up_lr_weight (deeper -> shallower, 深い層->浅い層): {up_lr_weight}")
else: else:
up_lr_weight = [1.0] * max_len_for_down_or_up
logger.info("up_lr_weight: all 1.0, すべて1.0") logger.info("up_lr_weight: all 1.0, すべて1.0")
return down_lr_weight, mid_lr_weight, up_lr_weight lr_weight = down_lr_weight + mid_lr_weight + up_lr_weight
if is_sdxl:
lr_weight = [1.0] + lr_weight + [1.0] # add 1.0 for emb_layers and out
assert (not is_sdxl and len(lr_weight) == LoRANetwork.NUM_OF_BLOCKS * 2 + LoRANetwork.NUM_OF_MID_BLOCKS) or (
is_sdxl and len(lr_weight) == 1 + LoRANetwork.SDXL_NUM_OF_BLOCKS * 2 + LoRANetwork.SDXL_NUM_OF_MID_BLOCKS + 1
), f"lr_weight length is invalid: {len(lr_weight)}"
return lr_weight
# lr_weightが0のblockをblock_dimsから除外する、外部から呼び出す可能性を考慮しておく # lr_weightが0のblockをblock_dimsから除外する、外部から呼び出す可能性を考慮しておく
def remove_block_dims_and_alphas( 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 is_sdxl, block_dims, block_alphas, conv_block_dims, conv_block_alphas, block_lr_weight: Optional[List[float]]
): ):
# set 0 to block dim without learning rate to remove the block if block_lr_weight is not None:
if down_lr_weight != None: for i, lr in enumerate(block_lr_weight):
for i, lr in enumerate(down_lr_weight):
if lr == 0: if lr == 0:
block_dims[i] = 0 block_dims[i] = 0
if conv_block_dims is not None: if conv_block_dims is not None:
conv_block_dims[i] = 0 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 return block_dims, block_alphas, conv_block_dims, conv_block_alphas
# 外部から呼び出す可能性を考慮しておく # 外部から呼び出す可能性を考慮しておく
def get_block_index(lora_name: str) -> int: def get_block_index(lora_name: str, is_sdxl: bool = False) -> int:
block_idx = -1 # invalid lora name block_idx = -1 # invalid lora name
if not is_sdxl:
m = RE_UPDOWN.search(lora_name) m = RE_UPDOWN.search(lora_name)
if m: if m:
g = m.groups() g = m.groups()
@@ -698,9 +734,22 @@ def get_block_index(lora_name: str) -> int:
block_idx = 1 + idx # 0に該当するLoRAは存在しない block_idx = 1 + idx # 0に該当するLoRAは存在しない
elif g[0] == "up": elif g[0] == "up":
block_idx = LoRANetwork.NUM_OF_BLOCKS + 1 + idx block_idx = LoRANetwork.NUM_OF_BLOCKS + 1 + idx
elif "mid_block_" in lora_name: elif "mid_block_" in lora_name:
block_idx = LoRANetwork.NUM_OF_BLOCKS # idx=12 block_idx = LoRANetwork.NUM_OF_BLOCKS # idx=12
else:
# copy from sdxl_train
if lora_name.startswith("lora_unet_"):
name = lora_name[len("lora_unet_") :]
if name.startswith("time_embed_") or name.startswith("label_emb_"): # No LoRA
block_idx = 0 # 0
elif name.startswith("input_blocks_"): # 1-9
block_idx = 1 + int(name.split("_")[2])
elif name.startswith("middle_block_"): # 10-12
block_idx = 10 + int(name.split("_")[2])
elif name.startswith("output_blocks_"): # 13-21
block_idx = 13 + int(name.split("_")[2])
elif name.startswith("out_"): # 22, out, no LoRA
block_idx = 22
return block_idx return block_idx
@@ -742,15 +791,18 @@ def create_network_from_weights(multiplier, file, vae, text_encoder, unet, weigh
) )
# block lr # block lr
down_lr_weight, mid_lr_weight, up_lr_weight = parse_block_lr_kwargs(kwargs) block_lr_weight = parse_block_lr_kwargs(kwargs)
if up_lr_weight is not None or mid_lr_weight is not None or down_lr_weight is not None: if block_lr_weight is not None:
network.set_block_lr_weight(up_lr_weight, mid_lr_weight, down_lr_weight) network.set_block_lr_weight(block_lr_weight)
return network, weights_sd return network, weights_sd
class LoRANetwork(torch.nn.Module): class LoRANetwork(torch.nn.Module):
NUM_OF_BLOCKS = 12 # フルモデル相当でのup,downの層の数 NUM_OF_BLOCKS = 12 # フルモデル相当でのup,downの層の数
NUM_OF_MID_BLOCKS = 1
SDXL_NUM_OF_BLOCKS = 9 # SDXLのモデルでのinput/outputの層の数 total=1(base) 9(input) + 3(mid) + 9(output) + 1(out) = 23
SDXL_NUM_OF_MID_BLOCKS = 3
UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel"] UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel"]
UNET_TARGET_REPLACE_MODULE_CONV2D_3X3 = ["ResnetBlock2D", "Downsample2D", "Upsample2D"] UNET_TARGET_REPLACE_MODULE_CONV2D_3X3 = ["ResnetBlock2D", "Downsample2D", "Upsample2D"]
@@ -782,6 +834,7 @@ class LoRANetwork(torch.nn.Module):
modules_alpha: Optional[Dict[str, int]] = None, modules_alpha: Optional[Dict[str, int]] = None,
module_class: Type[object] = LoRAModule, module_class: Type[object] = LoRAModule,
varbose: Optional[bool] = False, varbose: Optional[bool] = False,
is_sdxl: Optional[bool] = False,
) -> None: ) -> None:
""" """
LoRA network: すごく引数が多いが、パターンは以下の通り LoRA network: すごく引数が多いが、パターンは以下の通り
@@ -863,7 +916,7 @@ class LoRANetwork(torch.nn.Module):
alpha = modules_alpha[lora_name] alpha = modules_alpha[lora_name]
elif is_unet and block_dims is not None: elif is_unet and block_dims is not None:
# U-Netでblock_dims指定あり # U-Netでblock_dims指定あり
block_idx = get_block_index(lora_name) block_idx = get_block_index(lora_name, is_sdxl)
if is_linear or is_conv2d_1x1: if is_linear or is_conv2d_1x1:
dim = block_dims[block_idx] dim = block_dims[block_idx]
alpha = block_alphas[block_idx] alpha = block_alphas[block_idx]
@@ -927,15 +980,13 @@ class LoRANetwork(torch.nn.Module):
skipped = skipped_te + skipped_un skipped = skipped_te + skipped_un
if varbose and len(skipped) > 0: if varbose and len(skipped) > 0:
logger.warning( logger.warn(
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モジュールはスキップされます:" 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: for name in skipped:
logger.info(f"\t{name}") logger.info(f"\t{name}")
self.up_lr_weight: List[float] = None self.block_lr_weight = None
self.down_lr_weight: List[float] = None
self.mid_lr_weight: float = None
self.block_lr = False self.block_lr = False
# assertion # assertion
@@ -966,12 +1017,12 @@ class LoRANetwork(torch.nn.Module):
def apply_to(self, text_encoder, unet, apply_text_encoder=True, apply_unet=True): def apply_to(self, text_encoder, unet, apply_text_encoder=True, apply_unet=True):
if apply_text_encoder: if apply_text_encoder:
logger.info("enable LoRA for text encoder") logger.info(f"enable LoRA for text encoder: {len(self.text_encoder_loras)} modules")
else: else:
self.text_encoder_loras = [] self.text_encoder_loras = []
if apply_unet: if apply_unet:
logger.info("enable LoRA for U-Net") logger.info(f"enable LoRA for U-Net: {len(self.unet_loras)} modules")
else: else:
self.unet_loras = [] self.unet_loras = []
@@ -1012,34 +1063,14 @@ class LoRANetwork(torch.nn.Module):
logger.info(f"weights are merged") logger.info(f"weights are merged")
# 層別学習率用に層ごとの学習率に対する倍率を定義する　引数の順番が逆だがとりあえず気にしない # 層別学習率用に層ごとの学習率に対する倍率を定義する　引数の順番が逆だがとりあえず気にしない
def set_block_lr_weight( def set_block_lr_weight(self, block_lr_weight: Optional[List[float]]):
self,
up_lr_weight: List[float] = None,
mid_lr_weight: float = None,
down_lr_weight: List[float] = None,
):
self.block_lr = True self.block_lr = True
self.down_lr_weight = down_lr_weight self.block_lr_weight = block_lr_weight
self.mid_lr_weight = mid_lr_weight
self.up_lr_weight = up_lr_weight
def get_lr_weight(self, lora: LoRAModule) -> float: def get_lr_weight(self, block_idx: int) -> float:
lr_weight = 1.0 if not self.block_lr or self.block_lr_weight is None:
block_idx = get_block_index(lora.lora_name) return 1.0
if block_idx < 0: return self.block_lr_weight[block_idx]
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 set_loraplus_lr_ratio(self, loraplus_lr_ratio, loraplus_unet_lr_ratio, loraplus_text_encoder_lr_ratio): def set_loraplus_lr_ratio(self, loraplus_lr_ratio, loraplus_unet_lr_ratio, loraplus_text_encoder_lr_ratio):
self.loraplus_lr_ratio = loraplus_lr_ratio self.loraplus_lr_ratio = loraplus_lr_ratio
@@ -1106,10 +1137,16 @@ class LoRANetwork(torch.nn.Module):
if self.unet_loras: if self.unet_loras:
if self.block_lr: if self.block_lr:
is_sdxl = False
for lora in self.unet_loras:
if "input_blocks" in lora.lora_name or "output_blocks" in lora.lora_name:
is_sdxl = True
break
# 学習率のグラフをblockごとにしたいので、blockごとにloraを分類 # 学習率のグラフをblockごとにしたいので、blockごとにloraを分類
block_idx_to_lora = {} block_idx_to_lora = {}
for lora in self.unet_loras: for lora in self.unet_loras:
idx = get_block_index(lora.lora_name) idx = get_block_index(lora.lora_name, is_sdxl)
if idx not in block_idx_to_lora: if idx not in block_idx_to_lora:
block_idx_to_lora[idx] = [] block_idx_to_lora[idx] = []
block_idx_to_lora[idx].append(lora) block_idx_to_lora[idx].append(lora)
@@ -1118,7 +1155,7 @@ class LoRANetwork(torch.nn.Module):
for idx, block_loras in block_idx_to_lora.items(): for idx, block_loras in block_idx_to_lora.items():
params, descriptions = assemble_params( params, descriptions = assemble_params(
block_loras, block_loras,
(unet_lr if unet_lr is not None else default_lr) * self.get_lr_weight(block_loras[0]), (unet_lr if unet_lr is not None else default_lr) * self.get_lr_weight(idx),
self.loraplus_unet_lr_ratio or self.loraplus_lr_ratio, self.loraplus_unet_lr_ratio or self.loraplus_lr_ratio,
) )
all_params.extend(params) all_params.extend(params)

4
train_network.py
View File

@@ -346,13 +346,13 @@ class NetworkTrainer:
else: else:
trainable_params = results trainable_params = results
lr_descriptions = None lr_descriptions = None
except TypeError: except TypeError as e:
# logger.warning(f"{e}")
# accelerator.print( # accelerator.print(
# "Deprecated: use prepare_optimizer_params(text_encoder_lr, unet_lr, learning_rate) instead of prepare_optimizer_params(text_encoder_lr, unet_lr)" # "Deprecated: use prepare_optimizer_params(text_encoder_lr, unet_lr, learning_rate) instead of prepare_optimizer_params(text_encoder_lr, unet_lr)"
# ) # )
trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr) trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr)
lr_descriptions = None lr_descriptions = None
print(lr_descriptions)
optimizer_name, optimizer_args, optimizer = train_util.get_optimizer(args, trainable_params) optimizer_name, optimizer_args, optimizer = train_util.get_optimizer(args, trainable_params)