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

format by black

Browse Source
This commit is contained in:
Kohya S
2023-03-29 21:23:27 +09:00
parent b996f5a6d6
commit bf3674c1db
1 changed files with 312 additions and 287 deletions
Download Patch File Download Diff File Expand all files Collapse all files

91
networks/lora.py
View File

@@ -18,11 +18,11 @@ class LoRAModule(torch.nn.Module):
""" """
def __init__(self, lora_name, org_module: torch.nn.Module, multiplier=1.0, lora_dim=4, alpha=1): 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). """ """if alpha == 0 or None, alpha is rank (no scaling)."""
super().__init__() super().__init__()
self.lora_name = lora_name self.lora_name = lora_name
if org_module.__class__.__name__ == 'Conv2d': if org_module.__class__.__name__ == "Conv2d":
in_dim = org_module.in_channels in_dim = org_module.in_channels
out_dim = org_module.out_channels out_dim = org_module.out_channels
else: else:
@@ -36,7 +36,7 @@ class LoRAModule(torch.nn.Module):
# else: # else:
self.lora_dim = lora_dim self.lora_dim = lora_dim
if org_module.__class__.__name__ == 'Conv2d': if org_module.__class__.__name__ == "Conv2d":
kernel_size = org_module.kernel_size kernel_size = org_module.kernel_size
stride = org_module.stride stride = org_module.stride
padding = org_module.padding padding = org_module.padding
@@ -50,7 +50,7 @@ class LoRAModule(torch.nn.Module):
alpha = alpha.detach().float().numpy() # without casting, bf16 causes error alpha = alpha.detach().float().numpy() # without casting, bf16 causes error
alpha = self.lora_dim if alpha is None or alpha == 0 else alpha alpha = self.lora_dim if alpha is None or alpha == 0 else alpha
self.scale = alpha / self.lora_dim self.scale = alpha / self.lora_dim
self.register_buffer('alpha', torch.tensor(alpha)) # 定数として扱える self.register_buffer("alpha", torch.tensor(alpha)) # 定数として扱える
# same as microsoft's # same as microsoft's
torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5)) torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
@@ -87,7 +87,7 @@ class LoRAModule(torch.nn.Module):
else: else:
seq_len = x.size()[1] seq_len = x.size()[1]
ratio = math.sqrt((self.region.size()[0] * self.region.size()[1]) / seq_len) ratio = math.sqrt((self.region.size()[0] * self.region.size()[1]) / seq_len)
h = int(self.region.size()[0] / ratio + .5) h = int(self.region.size()[0] / ratio + 0.5)
w = seq_len // h w = seq_len // h
r = self.region.to(x.device) r = self.region.to(x.device)
@@ -95,7 +95,7 @@ class LoRAModule(torch.nn.Module):
r = r.to(torch.float) r = r.to(torch.float)
r = r.unsqueeze(0).unsqueeze(1) r = r.unsqueeze(0).unsqueeze(1)
# print(self.lora_name, self.region.size(), x.size(), r.size(), h, w) # print(self.lora_name, self.region.size(), x.size(), r.size(), h, w)
r = torch.nn.functional.interpolate(r, (h, w), mode='bilinear') r = torch.nn.functional.interpolate(r, (h, w), mode="bilinear")
r = r.to(x.dtype) r = r.to(x.dtype)
if len(x.size()) == 3: if len(x.size()) == 3:
@@ -111,8 +111,8 @@ def create_network(multiplier, network_dim, network_alpha, vae, text_encoder, un
network_dim = 4 # default network_dim = 4 # default
# extract dim/alpha for conv2d, and block dim # extract dim/alpha for conv2d, and block dim
conv_dim = kwargs.get('conv_dim', None) conv_dim = kwargs.get("conv_dim", None)
conv_alpha = kwargs.get('conv_alpha', None) conv_alpha = kwargs.get("conv_alpha", None)
if conv_dim is not None: if conv_dim is not None:
conv_dim = int(conv_dim) conv_dim = int(conv_dim)
if conv_alpha is None: if conv_alpha is None:
@@ -148,30 +148,38 @@ def create_network(multiplier, network_dim, network_alpha, vae, text_encoder, un
assert len(conv_block_alphas) == NUM_BLOCKS, f"Number of block alphas is not same to {NUM_BLOCKS}" assert len(conv_block_alphas) == NUM_BLOCKS, f"Number of block alphas is not same to {NUM_BLOCKS}"
""" """
network = LoRANetwork(text_encoder, unet, multiplier=multiplier, lora_dim=network_dim, network = LoRANetwork(
alpha=network_alpha, conv_lora_dim=conv_dim, conv_alpha=conv_alpha) text_encoder,
unet,
multiplier=multiplier,
lora_dim=network_dim,
alpha=network_alpha,
conv_lora_dim=conv_dim,
conv_alpha=conv_alpha,
)
return network return network
def create_network_from_weights(multiplier, file, vae, text_encoder, unet, weights_sd=None, **kwargs): def create_network_from_weights(multiplier, file, vae, text_encoder, unet, weights_sd=None, **kwargs):
if weights_sd is None: if weights_sd is None:
if os.path.splitext(file)[1] == '.safetensors': if os.path.splitext(file)[1] == ".safetensors":
from safetensors.torch import load_file, safe_open from safetensors.torch import load_file, safe_open
weights_sd = load_file(file) weights_sd = load_file(file)
else: else:
weights_sd = torch.load(file, map_location='cpu') weights_sd = torch.load(file, map_location="cpu")
# get dim/alpha mapping # get dim/alpha mapping
modules_dim = {} modules_dim = {}
modules_alpha = {} modules_alpha = {}
for key, value in weights_sd.items(): for key, value in weights_sd.items():
if '.' not in key: if "." not in key:
continue continue
lora_name = key.split('.')[0] lora_name = key.split(".")[0]
if 'alpha' in key: if "alpha" in key:
modules_alpha[lora_name] = value modules_alpha[lora_name] = value
elif 'lora_down' in key: elif "lora_down" in key:
dim = value.size()[0] dim = value.size()[0]
modules_dim[lora_name] = dim modules_dim[lora_name] = dim
# print(lora_name, value.size(), dim) # print(lora_name, value.size(), dim)
@@ -191,10 +199,21 @@ class LoRANetwork(torch.nn.Module):
UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"] UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"]
UNET_TARGET_REPLACE_MODULE_CONV2D_3X3 = ["ResnetBlock2D", "Downsample2D", "Upsample2D"] UNET_TARGET_REPLACE_MODULE_CONV2D_3X3 = ["ResnetBlock2D", "Downsample2D", "Upsample2D"]
TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"] TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
LORA_PREFIX_UNET = 'lora_unet' LORA_PREFIX_UNET = "lora_unet"
LORA_PREFIX_TEXT_ENCODER = 'lora_te' 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, modules_dim=None, modules_alpha=None) -> None: def __init__(
self,
text_encoder,
unet,
multiplier=1.0,
lora_dim=4,
alpha=1,
conv_lora_dim=None,
conv_alpha=None,
modules_dim=None,
modules_alpha=None,
) -> None:
super().__init__() super().__init__()
self.multiplier = multiplier self.multiplier = multiplier
@@ -225,8 +244,8 @@ class LoRANetwork(torch.nn.Module):
is_conv2d = child_module.__class__.__name__ == "Conv2d" is_conv2d = child_module.__class__.__name__ == "Conv2d"
is_conv2d_1x1 = is_conv2d and child_module.kernel_size == (1, 1) is_conv2d_1x1 = is_conv2d and child_module.kernel_size == (1, 1)
if is_linear or is_conv2d: if is_linear or is_conv2d:
lora_name = prefix + '.' + name + '.' + child_name lora_name = prefix + "." + name + "." + child_name
lora_name = lora_name.replace('.', '_') lora_name = lora_name.replace(".", "_")
if modules_dim is not None: if modules_dim is not None:
if lora_name not in modules_dim: if lora_name not in modules_dim:
@@ -247,8 +266,9 @@ class LoRANetwork(torch.nn.Module):
loras.append(lora) loras.append(lora)
return loras return loras
self.text_encoder_loras = create_modules(LoRANetwork.LORA_PREFIX_TEXT_ENCODER, self.text_encoder_loras = create_modules(
text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE) LoRANetwork.LORA_PREFIX_TEXT_ENCODER, text_encoder, LoRANetwork.TEXT_ENCODER_TARGET_REPLACE_MODULE
)
print(f"create LoRA for Text Encoder: {len(self.text_encoder_loras)} modules.") 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 # extend U-Net target modules if conv2d 3x3 is enabled, or load from weights
@@ -273,11 +293,12 @@ class LoRANetwork(torch.nn.Module):
lora.multiplier = self.multiplier lora.multiplier = self.multiplier
def load_weights(self, file): def load_weights(self, file):
if os.path.splitext(file)[1] == '.safetensors': if os.path.splitext(file)[1] == ".safetensors":
from safetensors.torch import load_file, safe_open from safetensors.torch import load_file, safe_open
self.weights_sd = load_file(file) self.weights_sd = load_file(file)
else: else:
self.weights_sd = torch.load(file, map_location='cpu') self.weights_sd = torch.load(file, map_location="cpu")
def apply_to(self, text_encoder, unet, apply_text_encoder=None, apply_unet=None): def apply_to(self, text_encoder, unet, apply_text_encoder=None, apply_unet=None):
if self.weights_sd: if self.weights_sd:
@@ -291,12 +312,16 @@ class LoRANetwork(torch.nn.Module):
if apply_text_encoder is None: if apply_text_encoder is None:
apply_text_encoder = weights_has_text_encoder apply_text_encoder = weights_has_text_encoder
else: else:
assert apply_text_encoder == weights_has_text_encoder, f"text encoder weights: {weights_has_text_encoder} but text encoder flag: {apply_text_encoder} / 重みとText Encoderのフラグが矛盾しています" assert (
apply_text_encoder == weights_has_text_encoder
), f"text encoder weights: {weights_has_text_encoder} but text encoder flag: {apply_text_encoder} / 重みとText Encoderのフラグが矛盾しています"
if apply_unet is None: if apply_unet is None:
apply_unet = weights_has_unet apply_unet = weights_has_unet
else: else:
assert apply_unet == weights_has_unet, f"u-net weights: {weights_has_unet} but u-net flag: {apply_unet} / 重みとU-Netのフラグが矛盾しています" assert (
apply_unet == weights_has_unet
), f"u-net weights: {weights_has_unet} but u-net flag: {apply_unet} / 重みとU-Netのフラグが矛盾しています"
else: else:
assert apply_text_encoder is not None and apply_unet is not None, f"internal error: flag not set" assert apply_text_encoder is not None and apply_unet is not None, f"internal error: flag not set"
@@ -334,15 +359,15 @@ class LoRANetwork(torch.nn.Module):
all_params = [] all_params = []
if self.text_encoder_loras: if self.text_encoder_loras:
param_data = {'params': enumerate_params(self.text_encoder_loras)} param_data = {"params": enumerate_params(self.text_encoder_loras)}
if text_encoder_lr is not None: if text_encoder_lr is not None:
param_data['lr'] = text_encoder_lr param_data["lr"] = text_encoder_lr
all_params.append(param_data) all_params.append(param_data)
if self.unet_loras: if self.unet_loras:
param_data = {'params': enumerate_params(self.unet_loras)} param_data = {"params": enumerate_params(self.unet_loras)}
if unet_lr is not None: if unet_lr is not None:
param_data['lr'] = unet_lr param_data["lr"] = unet_lr
all_params.append(param_data) all_params.append(param_data)
return all_params return all_params
@@ -368,7 +393,7 @@ class LoRANetwork(torch.nn.Module):
v = v.detach().clone().to("cpu").to(dtype) v = v.detach().clone().to("cpu").to(dtype)
state_dict[key] = v state_dict[key] = v
if os.path.splitext(file)[1] == '.safetensors': if os.path.splitext(file)[1] == ".safetensors":
from safetensors.torch import save_file from safetensors.torch import save_file
# Precalculate model hashes to save time on indexing # Precalculate model hashes to save time on indexing
@@ -382,7 +407,7 @@ class LoRANetwork(torch.nn.Module):
else: else:
torch.save(state_dict, file) torch.save(state_dict, file)
@ staticmethod @staticmethod
def set_regions(networks, image): def set_regions(networks, image):
image = image.astype(np.float32) / 255.0 image = image.astype(np.float32) / 255.0
for i, network in enumerate(networks[:3]): for i, network in enumerate(networks[:3]):