mirror of
https://github.com/kohya-ss/sd-scripts.git
synced 2026-04-09 06:45:09 +00:00
remove LoRA-ControlNet
This commit is contained in:
Kohya S
File diff suppressed because it is too large
Load Diff
@@ -1,444 +0,0 @@
|
|||||||
# 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
|
|
||||||
import torch
|
|
||||||
from diffusers import UNet2DConditionModel
|
|
||||||
|
|
||||||
from library import train_util
|
|
||||||
|
|
||||||
|
|
||||||
class ControlLoRAModule(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
|
|
||||||
self.lora_dim = lora_dim
|
|
||||||
|
|
||||||
if org_module.__class__.__name__ == 'Conv2d':
|
|
||||||
in_dim = org_module.in_channels
|
|
||||||
out_dim = org_module.out_channels
|
|
||||||
|
|
||||||
self.lora_dim = min(self.lora_dim, in_dim, out_dim)
|
|
||||||
if self.lora_dim != lora_dim:
|
|
||||||
print(f"{lora_name} dim (rank) is changed: {self.lora_dim}")
|
|
||||||
|
|
||||||
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:
|
|
||||||
in_dim = org_module.in_features
|
|
||||||
out_dim = org_module.out_features
|
|
||||||
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
|
|
||||||
|
|
||||||
def apply_to(self):
|
|
||||||
self.org_forward = self.org_module.forward
|
|
||||||
self.org_module.forward = self.forward
|
|
||||||
del self.org_module
|
|
||||||
|
|
||||||
def set_as_control_path(self, control_path):
|
|
||||||
self.is_control_path = control_path
|
|
||||||
|
|
||||||
def forward(self, x):
|
|
||||||
if not self.is_control_path:
|
|
||||||
return self.org_forward(x)
|
|
||||||
return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
|
|
||||||
|
|
||||||
|
|
||||||
class ControlLoRANetwork(torch.nn.Module):
|
|
||||||
# UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel", "Attention"]
|
|
||||||
# TEXT_ENCODER_TARGET_REPLACE_MODULE = ["CLIPAttention", "CLIPMLP"]
|
|
||||||
LORA_PREFIX_UNET = 'lora_unet'
|
|
||||||
LORA_PREFIX_TEXT_ENCODER = 'lora_te'
|
|
||||||
|
|
||||||
def __init__(self, unet, weights_sd, multiplier=1.0, lora_dim=4, alpha=1) -> None:
|
|
||||||
super().__init__()
|
|
||||||
self.multiplier = multiplier
|
|
||||||
self.lora_dim = lora_dim
|
|
||||||
self.alpha = alpha
|
|
||||||
|
|
||||||
# create module instances
|
|
||||||
def create_modules(prefix, root_module: torch.nn.Module) -> List[ControlLoRAModule]: # , target_replace_modules
|
|
||||||
loras = []
|
|
||||||
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():
|
|
||||||
if module.__class__.__name__ == "Linear" or module.__class__.__name__ == "Conv2d": # and module.kernel_size == (1, 1)):
|
|
||||||
lora_name = prefix + '.' + name # + '.' + child_name
|
|
||||||
lora_name = lora_name.replace('.', '_')
|
|
||||||
|
|
||||||
if weights_sd is None:
|
|
||||||
dim, alpha = self.lora_dim, self.alpha
|
|
||||||
else:
|
|
||||||
down_weight = weights_sd.get(lora_name + ".lora_down.weight", None)
|
|
||||||
if down_weight is None:
|
|
||||||
continue
|
|
||||||
dim = down_weight.size()[0]
|
|
||||||
alpha = weights_sd.get(lora_name + ".alpha", dim)
|
|
||||||
|
|
||||||
lora = ControlLoRAModule(lora_name, module, self.multiplier, dim, alpha)
|
|
||||||
loras.append(lora)
|
|
||||||
return loras
|
|
||||||
|
|
||||||
self.unet_loras = create_modules(ControlLoRANetwork.LORA_PREFIX_UNET, unet) # , LoRANetwork.UNET_TARGET_REPLACE_MODULE)
|
|
||||||
print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
|
|
||||||
|
|
||||||
# make control model
|
|
||||||
self.control_model = torch.nn.Module()
|
|
||||||
|
|
||||||
dims = [320, 320, 320, 320, 640, 640, 640, 1280, 1280, 1280, 1280, 1280]
|
|
||||||
zero_convs = torch.nn.ModuleList()
|
|
||||||
for i, dim in enumerate(dims):
|
|
||||||
sub_list = torch.nn.ModuleList([torch.nn.Conv2d(dim, dim, 1)])
|
|
||||||
zero_convs.append(sub_list)
|
|
||||||
self.control_model.add_module("zero_convs", zero_convs)
|
|
||||||
|
|
||||||
middle_block_out = torch.nn.Conv2d(1280, 1280, 1)
|
|
||||||
self.control_model.add_module("middle_block_out", torch.nn.ModuleList([middle_block_out]))
|
|
||||||
|
|
||||||
dims = [16, 16, 32, 32, 96, 96, 256, 320]
|
|
||||||
strides = [1, 1, 2, 1, 2, 1, 2, 1]
|
|
||||||
prev_dim = 3
|
|
||||||
input_hint_block = torch.nn.Sequential()
|
|
||||||
for i, (dim, stride) in enumerate(zip(dims, strides)):
|
|
||||||
input_hint_block.append(torch.nn.Conv2d(prev_dim, dim, 3, stride, 1))
|
|
||||||
if i < len(dims) - 1:
|
|
||||||
input_hint_block.append(torch.nn.SiLU())
|
|
||||||
prev_dim = dim
|
|
||||||
self.control_model.add_module("input_hint_block", input_hint_block)
|
|
||||||
|
|
||||||
|
|
||||||
# def load_weights(self, file):
|
|
||||||
# if os.path.splitext(file)[1] == '.safetensors':
|
|
||||||
# from safetensors.torch import load_file, safe_open
|
|
||||||
# self.weights_sd = load_file(file)
|
|
||||||
# else:
|
|
||||||
# self.weights_sd = torch.load(file, map_location='cpu')
|
|
||||||
|
|
||||||
def apply_to(self):
|
|
||||||
for lora in self.unet_loras:
|
|
||||||
lora.apply_to()
|
|
||||||
self.add_module(lora.lora_name, lora)
|
|
||||||
|
|
||||||
def call_unet(self, unet, hint, sample, timestep, encoder_hidden_states):
|
|
||||||
# control path
|
|
||||||
hint = hint.to(sample.dtype).to(sample.device)
|
|
||||||
guided_hint = self.control_model.input_hint_block(hint)
|
|
||||||
|
|
||||||
for lora_module in self.unet_loras:
|
|
||||||
lora_module.set_as_control_path(True)
|
|
||||||
|
|
||||||
outs = self.unet_forward(unet, guided_hint, None, sample, timestep, encoder_hidden_states)
|
|
||||||
|
|
||||||
# U-Net
|
|
||||||
for lora_module in self.unet_loras:
|
|
||||||
lora_module.set_as_control_path(False)
|
|
||||||
|
|
||||||
sample = self.unet_forward(unet, None, outs, sample, timestep, encoder_hidden_states)
|
|
||||||
|
|
||||||
return sample
|
|
||||||
|
|
||||||
def unet_forward(self, unet: UNet2DConditionModel, guided_hint, ctrl_outs, sample, timestep, encoder_hidden_states):
|
|
||||||
# copy from UNet2DConditionModel
|
|
||||||
default_overall_up_factor = 2**unet.num_upsamplers
|
|
||||||
|
|
||||||
forward_upsample_size = False
|
|
||||||
upsample_size = None
|
|
||||||
|
|
||||||
if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
|
|
||||||
print("Forward upsample size to force interpolation output size.")
|
|
||||||
forward_upsample_size = True
|
|
||||||
|
|
||||||
# 0. center input if necessary
|
|
||||||
if unet.config.center_input_sample:
|
|
||||||
sample = 2 * sample - 1.0
|
|
||||||
|
|
||||||
# 1. time
|
|
||||||
timesteps = timestep
|
|
||||||
if not torch.is_tensor(timesteps):
|
|
||||||
# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
|
|
||||||
# This would be a good case for the `match` statement (Python 3.10+)
|
|
||||||
is_mps = sample.device.type == "mps"
|
|
||||||
if isinstance(timestep, float):
|
|
||||||
dtype = torch.float32 if is_mps else torch.float64
|
|
||||||
else:
|
|
||||||
dtype = torch.int32 if is_mps else torch.int64
|
|
||||||
timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
|
|
||||||
elif len(timesteps.shape) == 0:
|
|
||||||
timesteps = timesteps[None].to(sample.device)
|
|
||||||
|
|
||||||
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
|
||||||
timesteps = timesteps.expand(sample.shape[0])
|
|
||||||
|
|
||||||
t_emb = unet.time_proj(timesteps)
|
|
||||||
|
|
||||||
# timesteps does not contain any weights and will always return f32 tensors
|
|
||||||
# but time_embedding might actually be running in fp16. so we need to cast here.
|
|
||||||
# there might be better ways to encapsulate this.
|
|
||||||
t_emb = t_emb.to(dtype=unet.dtype)
|
|
||||||
emb = unet.time_embedding(t_emb)
|
|
||||||
|
|
||||||
if ctrl_outs is None:
|
|
||||||
outs = [] # control path
|
|
||||||
|
|
||||||
# 2. pre-process
|
|
||||||
sample = unet.conv_in(sample)
|
|
||||||
if guided_hint is not None:
|
|
||||||
sample += guided_hint
|
|
||||||
if ctrl_outs is None:
|
|
||||||
outs.append(self.control_model.zero_convs[0][0](sample)) # , emb, encoder_hidden_states))
|
|
||||||
|
|
||||||
# 3. down
|
|
||||||
zc_idx = 1
|
|
||||||
down_block_res_samples = (sample,)
|
|
||||||
for downsample_block in unet.down_blocks:
|
|
||||||
if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
|
|
||||||
sample, res_samples = downsample_block(
|
|
||||||
hidden_states=sample,
|
|
||||||
temb=emb,
|
|
||||||
encoder_hidden_states=encoder_hidden_states,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
|
|
||||||
if ctrl_outs is None:
|
|
||||||
for rs in res_samples:
|
|
||||||
# print("zc", zc_idx, rs.size())
|
|
||||||
outs.append(self.control_model.zero_convs[zc_idx][0](rs)) # , emb, encoder_hidden_states))
|
|
||||||
zc_idx += 1
|
|
||||||
|
|
||||||
down_block_res_samples += res_samples
|
|
||||||
|
|
||||||
# 4. mid
|
|
||||||
sample = unet.mid_block(sample, emb, encoder_hidden_states=encoder_hidden_states)
|
|
||||||
if ctrl_outs is None:
|
|
||||||
outs.append(self.control_model.middle_block_out[0](sample))
|
|
||||||
return outs
|
|
||||||
if ctrl_outs is not None:
|
|
||||||
sample += ctrl_outs.pop()
|
|
||||||
|
|
||||||
# 5. up
|
|
||||||
for i, upsample_block in enumerate(unet.up_blocks):
|
|
||||||
is_final_block = i == len(unet.up_blocks) - 1
|
|
||||||
|
|
||||||
res_samples = down_block_res_samples[-len(upsample_block.resnets):]
|
|
||||||
down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
|
|
||||||
|
|
||||||
if ctrl_outs is not None and len(ctrl_outs) > 0:
|
|
||||||
res_samples = list(res_samples)
|
|
||||||
apply_ctrl_outs = ctrl_outs[-len(res_samples):]
|
|
||||||
ctrl_outs = ctrl_outs[:-len(res_samples)]
|
|
||||||
for j in range(len(res_samples)):
|
|
||||||
res_samples[j] = res_samples[j] + apply_ctrl_outs[j]
|
|
||||||
res_samples = tuple(res_samples)
|
|
||||||
|
|
||||||
# if we have not reached the final block and need to forward the
|
|
||||||
# upsample size, we do it here
|
|
||||||
if not is_final_block and forward_upsample_size:
|
|
||||||
upsample_size = down_block_res_samples[-1].shape[2:]
|
|
||||||
|
|
||||||
if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
|
|
||||||
sample = upsample_block(
|
|
||||||
hidden_states=sample,
|
|
||||||
temb=emb,
|
|
||||||
res_hidden_states_tuple=res_samples,
|
|
||||||
encoder_hidden_states=encoder_hidden_states,
|
|
||||||
upsample_size=upsample_size,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
sample = upsample_block(
|
|
||||||
hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size
|
|
||||||
)
|
|
||||||
# 6. post-process
|
|
||||||
sample = unet.conv_norm_out(sample)
|
|
||||||
sample = unet.conv_act(sample)
|
|
||||||
sample = unet.conv_out(sample)
|
|
||||||
|
|
||||||
return (sample,)
|
|
||||||
|
|
||||||
"""
|
|
||||||
default_overall_up_factor = 2**self.num_upsamplers
|
|
||||||
|
|
||||||
# upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
|
|
||||||
forward_upsample_size = False
|
|
||||||
upsample_size = None
|
|
||||||
|
|
||||||
if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
|
|
||||||
logger.info("Forward upsample size to force interpolation output size.")
|
|
||||||
forward_upsample_size = True
|
|
||||||
|
|
||||||
# 0. center input if necessary
|
|
||||||
if self.config.center_input_sample:
|
|
||||||
sample = 2 * sample - 1.0
|
|
||||||
|
|
||||||
# 1. time
|
|
||||||
timesteps = timestep
|
|
||||||
if not torch.is_tensor(timesteps):
|
|
||||||
# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
|
|
||||||
# This would be a good case for the `match` statement (Python 3.10+)
|
|
||||||
is_mps = sample.device.type == "mps"
|
|
||||||
if isinstance(timestep, float):
|
|
||||||
dtype = torch.float32 if is_mps else torch.float64
|
|
||||||
else:
|
|
||||||
dtype = torch.int32 if is_mps else torch.int64
|
|
||||||
timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
|
|
||||||
elif len(timesteps.shape) == 0:
|
|
||||||
timesteps = timesteps[None].to(sample.device)
|
|
||||||
|
|
||||||
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
|
||||||
timesteps = timesteps.expand(sample.shape[0])
|
|
||||||
|
|
||||||
t_emb = self.time_proj(timesteps)
|
|
||||||
|
|
||||||
# timesteps does not contain any weights and will always return f32 tensors
|
|
||||||
# but time_embedding might actually be running in fp16. so we need to cast here.
|
|
||||||
# there might be better ways to encapsulate this.
|
|
||||||
t_emb = t_emb.to(dtype=self.dtype)
|
|
||||||
emb = self.time_embedding(t_emb)
|
|
||||||
|
|
||||||
if self.config.num_class_embeds is not None:
|
|
||||||
if class_labels is None:
|
|
||||||
raise ValueError("class_labels should be provided when num_class_embeds > 0")
|
|
||||||
class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
|
|
||||||
emb = emb + class_emb
|
|
||||||
|
|
||||||
# 2. pre-process
|
|
||||||
sample = self.conv_in(sample)
|
|
||||||
|
|
||||||
# 3. down
|
|
||||||
down_block_res_samples = (sample,)
|
|
||||||
for downsample_block in self.down_blocks:
|
|
||||||
if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
|
|
||||||
sample, res_samples = downsample_block(
|
|
||||||
hidden_states=sample,
|
|
||||||
temb=emb,
|
|
||||||
encoder_hidden_states=encoder_hidden_states,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
|
|
||||||
|
|
||||||
down_block_res_samples += res_samples
|
|
||||||
|
|
||||||
# 4. mid
|
|
||||||
sample = self.mid_block(sample, emb, encoder_hidden_states=encoder_hidden_states)
|
|
||||||
|
|
||||||
# 5. up
|
|
||||||
for i, upsample_block in enumerate(self.up_blocks):
|
|
||||||
is_final_block = i == len(self.up_blocks) - 1
|
|
||||||
|
|
||||||
res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
|
|
||||||
down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
|
|
||||||
|
|
||||||
# if we have not reached the final block and need to forward the
|
|
||||||
# upsample size, we do it here
|
|
||||||
if not is_final_block and forward_upsample_size:
|
|
||||||
upsample_size = down_block_res_samples[-1].shape[2:]
|
|
||||||
|
|
||||||
if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
|
|
||||||
sample = upsample_block(
|
|
||||||
hidden_states=sample,
|
|
||||||
temb=emb,
|
|
||||||
res_hidden_states_tuple=res_samples,
|
|
||||||
encoder_hidden_states=encoder_hidden_states,
|
|
||||||
upsample_size=upsample_size,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
sample = upsample_block(
|
|
||||||
hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size
|
|
||||||
)
|
|
||||||
# 6. post-process
|
|
||||||
sample = self.conv_norm_out(sample)
|
|
||||||
sample = self.conv_act(sample)
|
|
||||||
sample = self.conv_out(sample)
|
|
||||||
|
|
||||||
if not return_dict:
|
|
||||||
return (sample,)
|
|
||||||
|
|
||||||
return UNet2DConditionOutput(sample=sample)
|
|
||||||
"""
|
|
||||||
|
|
||||||
def enable_gradient_checkpointing(self):
|
|
||||||
# not supported
|
|
||||||
pass
|
|
||||||
|
|
||||||
def prepare_optimizer_params(self, text_encoder_lr, unet_lr):
|
|
||||||
def enumerate_params(loras):
|
|
||||||
params = []
|
|
||||||
for lora in loras:
|
|
||||||
params.extend(lora.parameters())
|
|
||||||
return params
|
|
||||||
|
|
||||||
self.requires_grad_(True)
|
|
||||||
all_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:
|
|
||||||
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 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)
|
|
||||||
@@ -1,205 +0,0 @@
|
|||||||
# extract approximating LoRA by svd from SD 1.5 vs ControlNet
|
|
||||||
# https://github.com/lllyasviel/ControlNet/blob/main/tool_transfer_control.py
|
|
||||||
#
|
|
||||||
# The code is based on https://github.com/cloneofsimo/lora/blob/develop/lora_diffusion/cli_svd.py
|
|
||||||
# Thanks to cloneofsimo!
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import os
|
|
||||||
import torch
|
|
||||||
from safetensors.torch import load_file, save_file
|
|
||||||
from tqdm import tqdm
|
|
||||||
from diffusers import UNet2DConditionModel
|
|
||||||
|
|
||||||
import library.model_util as model_util
|
|
||||||
import control_net_lora
|
|
||||||
|
|
||||||
|
|
||||||
CLAMP_QUANTILE = 0.99
|
|
||||||
MIN_DIFF = 1e-6
|
|
||||||
|
|
||||||
|
|
||||||
def save_to_file(file_name, state_dict, dtype):
|
|
||||||
if dtype is not None:
|
|
||||||
for key in list(state_dict.keys()):
|
|
||||||
if type(state_dict[key]) == torch.Tensor:
|
|
||||||
state_dict[key] = state_dict[key].to(dtype)
|
|
||||||
|
|
||||||
if os.path.splitext(file_name)[1] == '.safetensors':
|
|
||||||
save_file(state_dict, file_name)
|
|
||||||
else:
|
|
||||||
torch.save(state_dict, file_name)
|
|
||||||
|
|
||||||
|
|
||||||
def svd(args):
|
|
||||||
def str_to_dtype(p):
|
|
||||||
if p == 'float':
|
|
||||||
return torch.float
|
|
||||||
if p == 'fp16':
|
|
||||||
return torch.float16
|
|
||||||
if p == 'bf16':
|
|
||||||
return torch.bfloat16
|
|
||||||
return None
|
|
||||||
|
|
||||||
save_dtype = str_to_dtype(args.save_precision)
|
|
||||||
|
|
||||||
# Diffusersのキーに変換するため、original sdとcontrol sdからU-Netに重みを読み込む ###############
|
|
||||||
|
|
||||||
# original sdをDiffusersのU-Netに読み込む
|
|
||||||
print(f"loading original SD model : {args.model_org}")
|
|
||||||
_, _, org_unet = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.model_org)
|
|
||||||
|
|
||||||
org_sd = torch.load(args.model_org, map_location='cpu')
|
|
||||||
if 'state_dict' in org_sd:
|
|
||||||
org_sd = org_sd['state_dict']
|
|
||||||
|
|
||||||
# control sdからキー変換しつつU-Netに対応する部分のみ取り出し、DiffusersのuU-Netに読み込む
|
|
||||||
print(f"loading control SD model : {args.model_tuned}")
|
|
||||||
|
|
||||||
ctrl_sd = torch.load(args.model_tuned, map_location='cpu')
|
|
||||||
ctrl_unet_sd = org_sd # あらかじめloadしておくことでcontrol sdにない部分はoriginal sdと同じにする
|
|
||||||
for key in list(ctrl_sd.keys()):
|
|
||||||
if key.startswith("control_"):
|
|
||||||
unet_key = "model.diffusion_" + key[len("control_"):]
|
|
||||||
if unet_key not in ctrl_unet_sd: # zero conv
|
|
||||||
continue
|
|
||||||
ctrl_unet_sd[unet_key] = ctrl_sd[key]
|
|
||||||
|
|
||||||
unet_config = model_util.create_unet_diffusers_config(args.v2)
|
|
||||||
ctrl_unet_sd_du = model_util.convert_ldm_unet_checkpoint(args.v2, ctrl_unet_sd, unet_config)
|
|
||||||
|
|
||||||
# load weights to U-Net
|
|
||||||
ctrl_unet = UNet2DConditionModel(**unet_config)
|
|
||||||
info = ctrl_unet.load_state_dict(ctrl_unet_sd_du)
|
|
||||||
print("loading control u-net:", info)
|
|
||||||
|
|
||||||
# LoRAに対応する部分のU-Netの重みを読み込む #################################
|
|
||||||
|
|
||||||
diffs = {}
|
|
||||||
for (org_name, org_module), (ctrl_name, ctrl_module) in zip(org_unet.named_modules(), ctrl_unet.named_modules()):
|
|
||||||
if org_module.__class__.__name__ != "Linear" and org_module.__class__.__name__ != "Conv2d":
|
|
||||||
continue
|
|
||||||
assert org_name == ctrl_name
|
|
||||||
|
|
||||||
lora_name = control_net_lora.ControlLoRANetwork.LORA_PREFIX_UNET + '.' + org_name # + '.' + child_name
|
|
||||||
lora_name = lora_name.replace('.', '_')
|
|
||||||
|
|
||||||
diff = ctrl_module.weight - org_module.weight
|
|
||||||
diff = diff.float()
|
|
||||||
|
|
||||||
if torch.max(torch.abs(diff)) < 1e-5:
|
|
||||||
# print(f"weights are same: {lora_name}")
|
|
||||||
continue
|
|
||||||
print(lora_name)
|
|
||||||
|
|
||||||
if args.device:
|
|
||||||
diff = diff.to(args.device)
|
|
||||||
|
|
||||||
diffs[lora_name] = diff
|
|
||||||
|
|
||||||
# make LoRA with svd
|
|
||||||
print("calculating by svd")
|
|
||||||
rank = args.dim
|
|
||||||
ctrl_lora_sd = {}
|
|
||||||
with torch.no_grad():
|
|
||||||
for lora_name, mat in tqdm(list(diffs.items())):
|
|
||||||
conv2d = (len(mat.size()) == 4)
|
|
||||||
kernel_size = None if not conv2d else mat.size()[2:]
|
|
||||||
|
|
||||||
if not conv2d or kernel_size == (1, 1):
|
|
||||||
if conv2d:
|
|
||||||
mat = mat.squeeze()
|
|
||||||
|
|
||||||
U, S, Vh = torch.linalg.svd(mat)
|
|
||||||
|
|
||||||
U = U[:, :rank]
|
|
||||||
S = S[:rank]
|
|
||||||
U = U @ torch.diag(S)
|
|
||||||
|
|
||||||
Vh = Vh[:rank, :]
|
|
||||||
|
|
||||||
dist = torch.cat([U.flatten(), Vh.flatten()])
|
|
||||||
hi_val = torch.quantile(dist, CLAMP_QUANTILE)
|
|
||||||
low_val = -hi_val
|
|
||||||
|
|
||||||
U = U.clamp(low_val, hi_val)
|
|
||||||
Vh = Vh.clamp(low_val, hi_val)
|
|
||||||
|
|
||||||
if conv2d:
|
|
||||||
U = U.unsqueeze(2).unsqueeze(3)
|
|
||||||
Vh = Vh.unsqueeze(2).unsqueeze(3)
|
|
||||||
else:
|
|
||||||
# conv2d kernel != (1,1)
|
|
||||||
in_channels = mat.size()[1]
|
|
||||||
current_rank = min(rank, in_channels, mat.size()[0])
|
|
||||||
if current_rank != rank:
|
|
||||||
print(f"channels of conv2d is too small. rank is changed to {current_rank} @ {lora_name}: {mat.size()}")
|
|
||||||
|
|
||||||
mat = mat.flatten(start_dim=1)
|
|
||||||
|
|
||||||
U, S, Vh = torch.linalg.svd(mat)
|
|
||||||
|
|
||||||
U = U[:, :current_rank]
|
|
||||||
S = S[:current_rank]
|
|
||||||
U = U @ torch.diag(S)
|
|
||||||
|
|
||||||
Vh = Vh[:current_rank, :]
|
|
||||||
|
|
||||||
dist = torch.cat([U.flatten(), Vh.flatten()])
|
|
||||||
hi_val = torch.quantile(dist, CLAMP_QUANTILE)
|
|
||||||
low_val = -hi_val
|
|
||||||
|
|
||||||
U = U.clamp(low_val, hi_val)
|
|
||||||
Vh = Vh.clamp(low_val, hi_val)
|
|
||||||
|
|
||||||
# U is (out_channels, rank) with 1x1 conv. So,
|
|
||||||
U = U.reshape(U.shape[0], U.shape[1], 1, 1)
|
|
||||||
# V is (rank, in_channels * kernel_size1 * kernel_size2)
|
|
||||||
# now reshape:
|
|
||||||
Vh = Vh.reshape(Vh.shape[0], in_channels, *kernel_size)
|
|
||||||
|
|
||||||
ctrl_lora_sd[lora_name + ".lora_up.weight"] = U
|
|
||||||
ctrl_lora_sd[lora_name + ".lora_down.weight"] = Vh
|
|
||||||
ctrl_lora_sd[lora_name + ".alpha"] = torch.tensor(current_rank)
|
|
||||||
|
|
||||||
# create LoRA from sd
|
|
||||||
lora_network = control_net_lora.ControlLoRANetwork(org_unet, ctrl_lora_sd, 1.0)
|
|
||||||
lora_network.apply_to()
|
|
||||||
|
|
||||||
for key, value in ctrl_sd.items():
|
|
||||||
if 'zero_convs' in key or 'input_hint_block' in key or 'middle_block_out' in key:
|
|
||||||
ctrl_lora_sd[key] = value
|
|
||||||
|
|
||||||
# verify state dict by loading it
|
|
||||||
info = lora_network.load_state_dict(ctrl_lora_sd)
|
|
||||||
print(f"loading control lora sd: {info}")
|
|
||||||
|
|
||||||
dir_name = os.path.dirname(args.save_to)
|
|
||||||
if dir_name and not os.path.exists(dir_name):
|
|
||||||
os.makedirs(dir_name, exist_ok=True)
|
|
||||||
|
|
||||||
# # minimum metadata
|
|
||||||
# metadata = {"ss_network_dim": str(args.dim), "ss_network_alpha": str(args.dim)}
|
|
||||||
|
|
||||||
# lora_network.save_weights(args.save_to, save_dtype, metadata)
|
|
||||||
save_to_file(args.save_to, ctrl_lora_sd, save_dtype)
|
|
||||||
print(f"LoRA weights are saved to: {args.save_to}")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
|
||||||
parser = argparse.ArgumentParser()
|
|
||||||
parser.add_argument("--v2", action='store_true',
|
|
||||||
help='load Stable Diffusion v2.x model / Stable Diffusion 2.xのモデルを読み込む')
|
|
||||||
parser.add_argument("--save_precision", type=str, default=None,
|
|
||||||
choices=[None, "float", "fp16", "bf16"], help="precision in saving, same to merging if omitted / 保存時に精度を変更して保存する、省略時はfloat")
|
|
||||||
parser.add_argument("--model_org", type=str, default=None,
|
|
||||||
help="Stable Diffusion original model: ckpt or safetensors file / 元モデル、ckptまたはsafetensors")
|
|
||||||
parser.add_argument("--model_tuned", type=str, default=None,
|
|
||||||
help="Stable Diffusion tuned model, LoRA is difference of `original to tuned`: ckpt or safetensors file / 派生モデル（生成されるLoRAは元→派生の差分になります）、ckptまたはsafetensors")
|
|
||||||
parser.add_argument("--save_to", type=str, default=None,
|
|
||||||
help="destination file name: ckpt or safetensors file / 保存先のファイル名、ckptまたはsafetensors")
|
|
||||||
parser.add_argument("--dim", type=int, default=4, help="dimension (rank) of LoRA (default 4) / LoRAの次元数（rank）（デフォルト4）")
|
|
||||||
parser.add_argument("--device", type=str, default=None, help="device to use, cuda for GPU / 計算を行うデバイス、cuda でGPUを使う")
|
|
||||||
|
|
||||||
args = parser.parse_args()
|
|
||||||
svd(args)
|
|
||||||
@@ -1,756 +0,0 @@
|
|||||||
from diffusers.optimization import SchedulerType, TYPE_TO_SCHEDULER_FUNCTION
|
|
||||||
from torch.optim import Optimizer
|
|
||||||
from typing import Optional, Tuple, Union
|
|
||||||
import importlib
|
|
||||||
import argparse
|
|
||||||
import gc
|
|
||||||
import math
|
|
||||||
import os
|
|
||||||
import random
|
|
||||||
import time
|
|
||||||
import json
|
|
||||||
|
|
||||||
from tqdm import tqdm
|
|
||||||
import torch
|
|
||||||
from accelerate.utils import set_seed
|
|
||||||
import diffusers
|
|
||||||
from diffusers import DDPMScheduler
|
|
||||||
|
|
||||||
import library.train_util as train_util
|
|
||||||
from library.train_util import BaseDataset, ImageInfo, glob_images
|
|
||||||
import networks.control_net_lora as control_net_rola
|
|
||||||
|
|
||||||
|
|
||||||
def collate_fn(examples):
|
|
||||||
return examples[0]
|
|
||||||
|
|
||||||
|
|
||||||
def generate_step_logs(args: argparse.Namespace, current_loss, avr_loss, lr_scheduler):
|
|
||||||
logs = {"loss/current": current_loss, "loss/average": avr_loss}
|
|
||||||
|
|
||||||
if args.network_train_unet_only:
|
|
||||||
logs["lr/unet"] = lr_scheduler.get_last_lr()[0]
|
|
||||||
elif args.network_train_text_encoder_only:
|
|
||||||
logs["lr/textencoder"] = lr_scheduler.get_last_lr()[0]
|
|
||||||
else:
|
|
||||||
logs["lr/textencoder"] = lr_scheduler.get_last_lr()[0]
|
|
||||||
logs["lr/unet"] = lr_scheduler.get_last_lr()[-1] # may be same to textencoder
|
|
||||||
|
|
||||||
return logs
|
|
||||||
|
|
||||||
|
|
||||||
# Monkeypatch newer get_scheduler() function overridng current version of diffusers.optimizer.get_scheduler
|
|
||||||
# code is taken from https://github.com/huggingface/diffusers diffusers.optimizer, commit d87cc15977b87160c30abaace3894e802ad9e1e6
|
|
||||||
# Which is a newer release of diffusers than currently packaged with sd-scripts
|
|
||||||
# This code can be removed when newer diffusers version (v0.12.1 or greater) is tested and implemented to sd-scripts
|
|
||||||
|
|
||||||
|
|
||||||
def get_scheduler_fix(
|
|
||||||
name: Union[str, SchedulerType],
|
|
||||||
optimizer: Optimizer,
|
|
||||||
num_warmup_steps: Optional[int] = None,
|
|
||||||
num_training_steps: Optional[int] = None,
|
|
||||||
num_cycles: int = 1,
|
|
||||||
power: float = 1.0,
|
|
||||||
):
|
|
||||||
"""
|
|
||||||
Unified API to get any scheduler from its name.
|
|
||||||
Args:
|
|
||||||
name (`str` or `SchedulerType`):
|
|
||||||
The name of the scheduler to use.
|
|
||||||
optimizer (`torch.optim.Optimizer`):
|
|
||||||
The optimizer that will be used during training.
|
|
||||||
num_warmup_steps (`int`, *optional*):
|
|
||||||
The number of warmup steps to do. This is not required by all schedulers (hence the argument being
|
|
||||||
optional), the function will raise an error if it's unset and the scheduler type requires it.
|
|
||||||
num_training_steps (`int``, *optional*):
|
|
||||||
The number of training steps to do. This is not required by all schedulers (hence the argument being
|
|
||||||
optional), the function will raise an error if it's unset and the scheduler type requires it.
|
|
||||||
num_cycles (`int`, *optional*):
|
|
||||||
The number of hard restarts used in `COSINE_WITH_RESTARTS` scheduler.
|
|
||||||
power (`float`, *optional*, defaults to 1.0):
|
|
||||||
Power factor. See `POLYNOMIAL` scheduler
|
|
||||||
last_epoch (`int`, *optional*, defaults to -1):
|
|
||||||
The index of the last epoch when resuming training.
|
|
||||||
"""
|
|
||||||
name = SchedulerType(name)
|
|
||||||
schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name]
|
|
||||||
if name == SchedulerType.CONSTANT:
|
|
||||||
return schedule_func(optimizer)
|
|
||||||
|
|
||||||
# All other schedulers require `num_warmup_steps`
|
|
||||||
if num_warmup_steps is None:
|
|
||||||
raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.")
|
|
||||||
|
|
||||||
if name == SchedulerType.CONSTANT_WITH_WARMUP:
|
|
||||||
return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)
|
|
||||||
|
|
||||||
# All other schedulers require `num_training_steps`
|
|
||||||
if num_training_steps is None:
|
|
||||||
raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.")
|
|
||||||
|
|
||||||
if name == SchedulerType.COSINE_WITH_RESTARTS:
|
|
||||||
return schedule_func(
|
|
||||||
optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, num_cycles=num_cycles
|
|
||||||
)
|
|
||||||
|
|
||||||
if name == SchedulerType.POLYNOMIAL:
|
|
||||||
return schedule_func(
|
|
||||||
optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, power=power
|
|
||||||
)
|
|
||||||
|
|
||||||
return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps)
|
|
||||||
|
|
||||||
|
|
||||||
class ImagesWithHintDataset(BaseDataset):
|
|
||||||
def __init__(self, batch_size, train_data_dir, tokenizer, max_token_length, caption_extension, shuffle_caption, shuffle_keep_tokens, resolution, enable_bucket, min_bucket_reso, max_bucket_reso, bucket_reso_steps, bucket_no_upscale, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset) -> None:
|
|
||||||
super().__init__(tokenizer, max_token_length, shuffle_caption, shuffle_keep_tokens,
|
|
||||||
resolution, flip_aug, color_aug, face_crop_aug_range, random_crop, debug_dataset)
|
|
||||||
|
|
||||||
assert resolution is not None, f"resolution is required / resolution（解像度）指定は必須です"
|
|
||||||
|
|
||||||
self.batch_size = batch_size
|
|
||||||
self.size = min(self.width, self.height) # 短いほう
|
|
||||||
self.latents_cache = None
|
|
||||||
|
|
||||||
self.enable_bucket = enable_bucket
|
|
||||||
if self.enable_bucket:
|
|
||||||
assert min(resolution) >= min_bucket_reso, f"min_bucket_reso must be equal or less than resolution / min_bucket_resoは最小解像度より大きくできません。解像度を大きくするかmin_bucket_resoを小さくしてください"
|
|
||||||
assert max(resolution) <= max_bucket_reso, f"max_bucket_reso must be equal or greater than resolution / max_bucket_resoは最大解像度より小さくできません。解像度を小さくするかmin_bucket_resoを大きくしてください"
|
|
||||||
self.min_bucket_reso = min_bucket_reso
|
|
||||||
self.max_bucket_reso = max_bucket_reso
|
|
||||||
self.bucket_reso_steps = bucket_reso_steps
|
|
||||||
self.bucket_no_upscale = bucket_no_upscale
|
|
||||||
else:
|
|
||||||
self.min_bucket_reso = None
|
|
||||||
self.max_bucket_reso = None
|
|
||||||
self.bucket_reso_steps = None # この情報は使われない
|
|
||||||
self.bucket_no_upscale = False
|
|
||||||
|
|
||||||
# fill50k
|
|
||||||
print("loading fill50k dataset")
|
|
||||||
with open(os.path.join(train_data_dir, "prompt.json")) as f:
|
|
||||||
annos = f.readlines()
|
|
||||||
|
|
||||||
captions = []
|
|
||||||
src_paths = []
|
|
||||||
trg_paths = []
|
|
||||||
for anno in annos:
|
|
||||||
anno1 = json.loads(anno)
|
|
||||||
captions.append(anno1["prompt"])
|
|
||||||
src_paths.append(os.path.join(train_data_dir, anno1["source"]))
|
|
||||||
trg_paths.append(os.path.join(train_data_dir, anno1["target"]))
|
|
||||||
|
|
||||||
self.set_tag_frequency(os.path.basename(train_data_dir), captions) # タグ頻度を記録
|
|
||||||
self.dataset_dirs_info[os.path.basename(train_data_dir)] = {"n_repeats": 1, "img_count": len(src_paths)}
|
|
||||||
|
|
||||||
for src_path, trg_path, caption in zip(src_paths, trg_paths, captions):
|
|
||||||
info = ImageInfo(src_path, 1, caption, False, src_path)
|
|
||||||
self.register_image(info)
|
|
||||||
|
|
||||||
num_train_images = len(src_paths)
|
|
||||||
print(f"{num_train_images} train images with repeating.")
|
|
||||||
self.num_train_images = num_train_images
|
|
||||||
self.num_reg_images = 0
|
|
||||||
|
|
||||||
"""
|
|
||||||
def read_caption(img_path):
|
|
||||||
# captionの候補ファイル名を作る
|
|
||||||
base_name = os.path.splitext(img_path)[0]
|
|
||||||
base_name_face_det = base_name
|
|
||||||
tokens = base_name.split("_")
|
|
||||||
if len(tokens) >= 5:
|
|
||||||
base_name_face_det = "_".join(tokens[:-4])
|
|
||||||
cap_paths = [base_name + caption_extension, base_name_face_det + caption_extension]
|
|
||||||
|
|
||||||
caption = None
|
|
||||||
for cap_path in cap_paths:
|
|
||||||
if os.path.isfile(cap_path):
|
|
||||||
with open(cap_path, "rt", encoding='utf-8') as f:
|
|
||||||
try:
|
|
||||||
lines = f.readlines()
|
|
||||||
except UnicodeDecodeError as e:
|
|
||||||
print(f"illegal char in file (not UTF-8) / ファイルにUTF-8以外の文字があります: {cap_path}")
|
|
||||||
raise e
|
|
||||||
assert len(lines) > 0, f"caption file is empty / キャプションファイルが空です: {cap_path}"
|
|
||||||
caption = lines[0].strip()
|
|
||||||
break
|
|
||||||
return caption
|
|
||||||
|
|
||||||
def load_dreambooth_dir(dir):
|
|
||||||
if not os.path.isdir(dir):
|
|
||||||
# print(f"ignore file: {dir}")
|
|
||||||
return 0, [], []
|
|
||||||
|
|
||||||
tokens = os.path.basename(dir).split('_')
|
|
||||||
try:
|
|
||||||
n_repeats = int(tokens[0])
|
|
||||||
except ValueError as e:
|
|
||||||
print(f"ignore directory without repeats / 繰り返し回数のないディレクトリを無視します: {dir}")
|
|
||||||
return 0, [], []
|
|
||||||
|
|
||||||
caption_by_folder = '_'.join(tokens[1:])
|
|
||||||
img_paths = glob_images(dir, "*")
|
|
||||||
print(f"found directory {n_repeats}_{caption_by_folder} contains {len(img_paths)} image files")
|
|
||||||
|
|
||||||
# 画像ファイルごとにプロンプトを読み込み、もしあればそちらを使う
|
|
||||||
captions = []
|
|
||||||
for img_path in img_paths:
|
|
||||||
cap_for_img = read_caption(img_path)
|
|
||||||
captions.append(caption_by_folder if cap_for_img is None else cap_for_img)
|
|
||||||
|
|
||||||
self.set_tag_frequency(os.path.basename(dir), captions) # タグ頻度を記録
|
|
||||||
|
|
||||||
return n_repeats, img_paths, captions
|
|
||||||
|
|
||||||
print("prepare train images.")
|
|
||||||
train_dirs = os.listdir(train_data_dir)
|
|
||||||
num_train_images = 0
|
|
||||||
for dir in train_dirs:
|
|
||||||
n_repeats, img_paths, captions = load_dreambooth_dir(os.path.join(train_data_dir, dir))
|
|
||||||
num_train_images += n_repeats * len(img_paths)
|
|
||||||
|
|
||||||
for img_path, caption in zip(img_paths, captions):
|
|
||||||
info = ImageInfo(img_path, n_repeats, caption, False, img_path)
|
|
||||||
self.register_image(info)
|
|
||||||
|
|
||||||
self.dataset_dirs_info[os.path.basename(dir)] = {"n_repeats": n_repeats, "img_count": len(img_paths)}
|
|
||||||
|
|
||||||
print(f"{num_train_images} train images with repeating.")
|
|
||||||
self.num_train_images = num_train_images
|
|
||||||
self.num_reg_images = 0
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __getitem__(self, index):
|
|
||||||
# latentsのcacheをサポートしてない
|
|
||||||
if index == 0:
|
|
||||||
self.shuffle_buckets()
|
|
||||||
|
|
||||||
bucket = self.bucket_manager.buckets[self.buckets_indices[index].bucket_index]
|
|
||||||
bucket_batch_size = self.buckets_indices[index].bucket_batch_size
|
|
||||||
image_index = self.buckets_indices[index].batch_index * bucket_batch_size
|
|
||||||
|
|
||||||
loss_weights = []
|
|
||||||
captions = []
|
|
||||||
input_ids_list = []
|
|
||||||
images = []
|
|
||||||
hint_images = []
|
|
||||||
|
|
||||||
for image_key in bucket[image_index:image_index + bucket_batch_size]:
|
|
||||||
image_info = self.image_data[image_key]
|
|
||||||
loss_weights.append(1.0)
|
|
||||||
|
|
||||||
# image/latentsを処理する
|
|
||||||
# 画像を読み込み、必要ならcropする
|
|
||||||
src_path = image_info.absolute_path
|
|
||||||
trg_path = src_path.replace("source", "target")
|
|
||||||
img, face_cx, face_cy, face_w, face_h = self.load_image_with_face_info(trg_path)
|
|
||||||
hint_img, face_cx, face_cy, face_w, face_h = self.load_image_with_face_info(src_path)
|
|
||||||
assert img.shape[0:2] == hint_img.shape[0:2]
|
|
||||||
im_h, im_w = img.shape[0:2]
|
|
||||||
|
|
||||||
if self.enable_bucket:
|
|
||||||
img = self.trim_and_resize_if_required(img, image_info.bucket_reso, image_info.resized_size)
|
|
||||||
else:
|
|
||||||
if face_cx > 0: # 顔位置情報あり
|
|
||||||
img = self.crop_target(img, face_cx, face_cy, face_w, face_h)
|
|
||||||
elif im_h > self.height or im_w > self.width:
|
|
||||||
assert self.random_crop, f"image too large, but cropping and bucketing are disabled / 画像サイズが大きいのでface_crop_aug_rangeかrandom_crop、またはbucketを有効にしてください: {image_info.absolute_path}"
|
|
||||||
if im_h > self.height:
|
|
||||||
p = random.randint(0, im_h - self.height)
|
|
||||||
img = img[p:p + self.height]
|
|
||||||
if im_w > self.width:
|
|
||||||
p = random.randint(0, im_w - self.width)
|
|
||||||
img = img[:, p:p + self.width]
|
|
||||||
|
|
||||||
im_h, im_w = img.shape[0:2]
|
|
||||||
assert im_h == self.height and im_w == self.width, f"image size is small / 画像サイズが小さいようです: {image_info.absolute_path}"
|
|
||||||
|
|
||||||
# augmentation
|
|
||||||
if self.aug is not None:
|
|
||||||
# TODO color aug does not work
|
|
||||||
auged = self.aug(image=img, image2=hint_img)
|
|
||||||
img = auged['image']
|
|
||||||
hint_img = auged['image2']
|
|
||||||
|
|
||||||
image = self.image_transforms(img) # -1.0~1.0のtorch.Tensorになる
|
|
||||||
hint_image = self.image_transforms(hint_img) # -1.0~1.0のtorch.Tensorになる
|
|
||||||
|
|
||||||
images.append(image)
|
|
||||||
hint_images.append(hint_image)
|
|
||||||
|
|
||||||
caption = self.process_caption(image_info.caption)
|
|
||||||
captions.append(caption)
|
|
||||||
if not self.token_padding_disabled: # this option might be omitted in future
|
|
||||||
input_ids_list.append(self.get_input_ids(caption))
|
|
||||||
|
|
||||||
example = {}
|
|
||||||
example['loss_weights'] = torch.FloatTensor(loss_weights)
|
|
||||||
|
|
||||||
if self.token_padding_disabled:
|
|
||||||
# padding=True means pad in the batch
|
|
||||||
example['input_ids'] = self.tokenizer(captions, padding=True, truncation=True, return_tensors="pt").input_ids
|
|
||||||
else:
|
|
||||||
# batch processing seems to be good
|
|
||||||
example['input_ids'] = torch.stack(input_ids_list)
|
|
||||||
|
|
||||||
images = torch.stack(images)
|
|
||||||
images = images.to(memory_format=torch.contiguous_format).float()
|
|
||||||
example['images'] = images
|
|
||||||
|
|
||||||
hint_images = torch.stack(hint_images)
|
|
||||||
hint_images = hint_images.to(memory_format=torch.contiguous_format).float()
|
|
||||||
example['hint_images'] = hint_images
|
|
||||||
|
|
||||||
example['latents'] = None
|
|
||||||
|
|
||||||
if self.debug_dataset:
|
|
||||||
example['image_keys'] = bucket[image_index:image_index + self.batch_size]
|
|
||||||
example['captions'] = captions
|
|
||||||
return example
|
|
||||||
|
|
||||||
|
|
||||||
def train(args):
|
|
||||||
session_id = random.randint(0, 2**32)
|
|
||||||
training_started_at = time.time()
|
|
||||||
train_util.verify_training_args(args)
|
|
||||||
train_util.prepare_dataset_args(args, True)
|
|
||||||
|
|
||||||
cache_latents = args.cache_latents
|
|
||||||
use_dreambooth_method = args.in_json is None
|
|
||||||
|
|
||||||
if args.seed is not None:
|
|
||||||
set_seed(args.seed)
|
|
||||||
|
|
||||||
tokenizer = train_util.load_tokenizer(args)
|
|
||||||
|
|
||||||
# データセットを準備する
|
|
||||||
train_dataset = ImagesWithHintDataset(args.train_batch_size, args.train_data_dir,
|
|
||||||
tokenizer, args.max_token_length, args.caption_extension, args.shuffle_caption, args.keep_tokens,
|
|
||||||
args.resolution, args.enable_bucket, args.min_bucket_reso, args.max_bucket_reso,
|
|
||||||
args.bucket_reso_steps, args.bucket_no_upscale,
|
|
||||||
args.flip_aug, args.color_aug, args.face_crop_aug_range,
|
|
||||||
args.random_crop, args.debug_dataset)
|
|
||||||
|
|
||||||
# 学習データのdropout率を設定する
|
|
||||||
train_dataset.set_caption_dropout(args.caption_dropout_rate, args.caption_dropout_every_n_epochs, args.caption_tag_dropout_rate)
|
|
||||||
|
|
||||||
train_dataset.make_buckets()
|
|
||||||
|
|
||||||
if args.debug_dataset:
|
|
||||||
train_util.debug_dataset(train_dataset)
|
|
||||||
return
|
|
||||||
if len(train_dataset) == 0:
|
|
||||||
print("No data found. Please verify arguments (train_data_dir must be the parent of folders with images) / 画像がありません。引数指定を確認してください（train_data_dirには画像があるフォルダではなく、画像があるフォルダの親フォルダを指定する必要があります）")
|
|
||||||
return
|
|
||||||
|
|
||||||
# acceleratorを準備する
|
|
||||||
print("prepare accelerator")
|
|
||||||
accelerator, unwrap_model = train_util.prepare_accelerator(args)
|
|
||||||
|
|
||||||
# mixed precisionに対応した型を用意しておき適宜castする
|
|
||||||
weight_dtype, save_dtype = train_util.prepare_dtype(args)
|
|
||||||
|
|
||||||
# モデルを読み込む
|
|
||||||
text_encoder, vae, unet, _ = train_util.load_target_model(args, weight_dtype)
|
|
||||||
|
|
||||||
# モデルに xformers とか memory efficient attention を組み込む
|
|
||||||
train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers)
|
|
||||||
|
|
||||||
# 学習を準備する
|
|
||||||
if cache_latents:
|
|
||||||
vae.to(accelerator.device, dtype=weight_dtype)
|
|
||||||
vae.requires_grad_(False)
|
|
||||||
vae.eval()
|
|
||||||
with torch.no_grad():
|
|
||||||
train_dataset.cache_latents(vae)
|
|
||||||
vae.to("cpu")
|
|
||||||
if torch.cuda.is_available():
|
|
||||||
torch.cuda.empty_cache()
|
|
||||||
gc.collect()
|
|
||||||
|
|
||||||
# prepare network
|
|
||||||
print("import network module:", args.network_module)
|
|
||||||
network_module = importlib.import_module(args.network_module)
|
|
||||||
|
|
||||||
net_kwargs = {}
|
|
||||||
if args.network_args is not None:
|
|
||||||
for net_arg in args.network_args:
|
|
||||||
key, value = net_arg.split('=')
|
|
||||||
net_kwargs[key] = value
|
|
||||||
|
|
||||||
# if a new network is added in future, add if ~ then blocks for each network (;'∀')
|
|
||||||
network: control_net_rola.ControlLoRANetwork = network_module.create_network(
|
|
||||||
1.0, args.network_dim, args.network_alpha, vae, text_encoder, unet, **net_kwargs)
|
|
||||||
if network is None:
|
|
||||||
return
|
|
||||||
|
|
||||||
if args.network_weights is not None:
|
|
||||||
print("load network weights from:", args.network_weights)
|
|
||||||
network.load_weights(args.network_weights)
|
|
||||||
|
|
||||||
train_unet = not args.network_train_text_encoder_only
|
|
||||||
train_text_encoder = not args.network_train_unet_only
|
|
||||||
network.apply_to(text_encoder, unet, train_text_encoder, train_unet)
|
|
||||||
|
|
||||||
if args.gradient_checkpointing:
|
|
||||||
unet.enable_gradient_checkpointing()
|
|
||||||
text_encoder.gradient_checkpointing_enable()
|
|
||||||
network.enable_gradient_checkpointing() # may have no effect
|
|
||||||
|
|
||||||
# 学習に必要なクラスを準備する
|
|
||||||
print("prepare optimizer, data loader etc.")
|
|
||||||
|
|
||||||
# 8-bit Adamを使う
|
|
||||||
if args.use_8bit_adam:
|
|
||||||
try:
|
|
||||||
import bitsandbytes as bnb
|
|
||||||
except ImportError:
|
|
||||||
raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
|
|
||||||
print("use 8-bit Adam optimizer")
|
|
||||||
optimizer_class = bnb.optim.AdamW8bit
|
|
||||||
else:
|
|
||||||
optimizer_class = torch.optim.AdamW
|
|
||||||
|
|
||||||
trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr)
|
|
||||||
|
|
||||||
# betaやweight decayはdiffusers DreamBoothもDreamBooth SDもデフォルト値のようなのでオプションはとりあえず省略
|
|
||||||
optimizer = optimizer_class(trainable_params, lr=args.learning_rate)
|
|
||||||
|
|
||||||
# dataloaderを準備する
|
|
||||||
# DataLoaderのプロセス数：0はメインプロセスになる
|
|
||||||
n_workers = min(args.max_data_loader_n_workers, os.cpu_count() - 1) # cpu_count-1 ただし最大で指定された数まで
|
|
||||||
train_dataloader = torch.utils.data.DataLoader(
|
|
||||||
train_dataset, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=n_workers, persistent_workers=args.persistent_data_loader_workers)
|
|
||||||
|
|
||||||
# 学習ステップ数を計算する
|
|
||||||
if args.max_train_epochs is not None:
|
|
||||||
args.max_train_steps = args.max_train_epochs * len(train_dataloader)
|
|
||||||
print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")
|
|
||||||
|
|
||||||
# lr schedulerを用意する
|
|
||||||
# lr_scheduler = diffusers.optimization.get_scheduler(
|
|
||||||
lr_scheduler = get_scheduler_fix(
|
|
||||||
args.lr_scheduler, optimizer, num_warmup_steps=args.lr_warmup_steps,
|
|
||||||
num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
|
|
||||||
num_cycles=args.lr_scheduler_num_cycles, power=args.lr_scheduler_power)
|
|
||||||
|
|
||||||
# 実験的機能：勾配も含めたfp16学習を行う モデル全体をfp16にする
|
|
||||||
if args.full_fp16:
|
|
||||||
assert args.mixed_precision == "fp16", "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
|
|
||||||
print("enable full fp16 training.")
|
|
||||||
network.to(weight_dtype)
|
|
||||||
|
|
||||||
# acceleratorがなんかよろしくやってくれるらしい
|
|
||||||
if train_unet and train_text_encoder:
|
|
||||||
unet, text_encoder, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
|
|
||||||
unet, text_encoder, network, optimizer, train_dataloader, lr_scheduler)
|
|
||||||
elif train_unet:
|
|
||||||
unet, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
|
|
||||||
unet, network, optimizer, train_dataloader, lr_scheduler)
|
|
||||||
elif train_text_encoder:
|
|
||||||
text_encoder, network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
|
|
||||||
text_encoder, network, optimizer, train_dataloader, lr_scheduler)
|
|
||||||
else:
|
|
||||||
network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
|
|
||||||
network, optimizer, train_dataloader, lr_scheduler)
|
|
||||||
|
|
||||||
unet.requires_grad_(False)
|
|
||||||
unet.to(accelerator.device, dtype=weight_dtype)
|
|
||||||
text_encoder.requires_grad_(False)
|
|
||||||
text_encoder.to(accelerator.device, dtype=weight_dtype)
|
|
||||||
if args.gradient_checkpointing: # according to TI example in Diffusers, train is required
|
|
||||||
unet.train()
|
|
||||||
text_encoder.train()
|
|
||||||
|
|
||||||
# set top parameter requires_grad = True for gradient checkpointing works
|
|
||||||
text_encoder.text_model.embeddings.requires_grad_(True)
|
|
||||||
else:
|
|
||||||
unet.eval()
|
|
||||||
text_encoder.eval()
|
|
||||||
|
|
||||||
network.prepare_grad_etc(text_encoder, unet)
|
|
||||||
|
|
||||||
if not cache_latents:
|
|
||||||
vae.requires_grad_(False)
|
|
||||||
vae.eval()
|
|
||||||
vae.to(accelerator.device, dtype=weight_dtype)
|
|
||||||
|
|
||||||
# 実験的機能：勾配も含めたfp16学習を行う PyTorchにパッチを当ててfp16でのgrad scaleを有効にする
|
|
||||||
if args.full_fp16:
|
|
||||||
train_util.patch_accelerator_for_fp16_training(accelerator)
|
|
||||||
|
|
||||||
# resumeする
|
|
||||||
if args.resume is not None:
|
|
||||||
print(f"resume training from state: {args.resume}")
|
|
||||||
accelerator.load_state(args.resume)
|
|
||||||
|
|
||||||
# epoch数を計算する
|
|
||||||
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
|
|
||||||
num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
|
|
||||||
if (args.save_n_epoch_ratio is not None) and (args.save_n_epoch_ratio > 0):
|
|
||||||
args.save_every_n_epochs = math.floor(num_train_epochs / args.save_n_epoch_ratio) or 1
|
|
||||||
|
|
||||||
# 学習する
|
|
||||||
total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
|
|
||||||
print("running training / 学習開始")
|
|
||||||
print(f" num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset.num_train_images}")
|
|
||||||
print(f" num reg images / 正則化画像の数: {train_dataset.num_reg_images}")
|
|
||||||
print(f" num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")
|
|
||||||
print(f" num epochs / epoch数: {num_train_epochs}")
|
|
||||||
print(f" batch size per device / バッチサイズ: {args.train_batch_size}")
|
|
||||||
print(f" total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")
|
|
||||||
print(f" gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")
|
|
||||||
print(f" total optimization steps / 学習ステップ数: {args.max_train_steps}")
|
|
||||||
|
|
||||||
metadata = {
|
|
||||||
"ss_session_id": session_id, # random integer indicating which group of epochs the model came from
|
|
||||||
"ss_training_started_at": training_started_at, # unix timestamp
|
|
||||||
"ss_output_name": args.output_name,
|
|
||||||
"ss_learning_rate": args.learning_rate,
|
|
||||||
"ss_text_encoder_lr": args.text_encoder_lr,
|
|
||||||
"ss_unet_lr": args.unet_lr,
|
|
||||||
"ss_num_train_images": train_dataset.num_train_images, # includes repeating
|
|
||||||
"ss_num_reg_images": train_dataset.num_reg_images,
|
|
||||||
"ss_num_batches_per_epoch": len(train_dataloader),
|
|
||||||
"ss_num_epochs": num_train_epochs,
|
|
||||||
"ss_batch_size_per_device": args.train_batch_size,
|
|
||||||
"ss_total_batch_size": total_batch_size,
|
|
||||||
"ss_gradient_checkpointing": args.gradient_checkpointing,
|
|
||||||
"ss_gradient_accumulation_steps": args.gradient_accumulation_steps,
|
|
||||||
"ss_max_train_steps": args.max_train_steps,
|
|
||||||
"ss_lr_warmup_steps": args.lr_warmup_steps,
|
|
||||||
"ss_lr_scheduler": args.lr_scheduler,
|
|
||||||
"ss_network_module": "control_net_" + args.network_module,
|
|
||||||
"ss_network_dim": args.network_dim, # None means default because another network than LoRA may have another default dim
|
|
||||||
"ss_network_alpha": args.network_alpha, # some networks may not use this value
|
|
||||||
"ss_mixed_precision": args.mixed_precision,
|
|
||||||
"ss_full_fp16": bool(args.full_fp16),
|
|
||||||
"ss_v2": bool(args.v2),
|
|
||||||
"ss_resolution": args.resolution,
|
|
||||||
"ss_clip_skip": args.clip_skip,
|
|
||||||
"ss_max_token_length": args.max_token_length,
|
|
||||||
"ss_color_aug": bool(args.color_aug),
|
|
||||||
"ss_flip_aug": bool(args.flip_aug),
|
|
||||||
"ss_random_crop": bool(args.random_crop),
|
|
||||||
"ss_shuffle_caption": bool(args.shuffle_caption),
|
|
||||||
"ss_cache_latents": bool(args.cache_latents),
|
|
||||||
"ss_enable_bucket": bool(train_dataset.enable_bucket),
|
|
||||||
"ss_min_bucket_reso": train_dataset.min_bucket_reso,
|
|
||||||
"ss_max_bucket_reso": train_dataset.max_bucket_reso,
|
|
||||||
"ss_seed": args.seed,
|
|
||||||
"ss_keep_tokens": args.keep_tokens,
|
|
||||||
"ss_dataset_dirs": json.dumps(train_dataset.dataset_dirs_info),
|
|
||||||
"ss_reg_dataset_dirs": json.dumps(train_dataset.reg_dataset_dirs_info),
|
|
||||||
"ss_tag_frequency": json.dumps(train_dataset.tag_frequency),
|
|
||||||
"ss_bucket_info": json.dumps(train_dataset.bucket_info),
|
|
||||||
"ss_training_comment": args.training_comment # will not be updated after training
|
|
||||||
}
|
|
||||||
|
|
||||||
# uncomment if another network is added
|
|
||||||
# for key, value in net_kwargs.items():
|
|
||||||
# metadata["ss_arg_" + key] = value
|
|
||||||
|
|
||||||
if args.pretrained_model_name_or_path is not None:
|
|
||||||
sd_model_name = args.pretrained_model_name_or_path
|
|
||||||
if os.path.exists(sd_model_name):
|
|
||||||
metadata["ss_sd_model_hash"] = train_util.model_hash(sd_model_name)
|
|
||||||
metadata["ss_new_sd_model_hash"] = train_util.calculate_sha256(sd_model_name)
|
|
||||||
sd_model_name = os.path.basename(sd_model_name)
|
|
||||||
metadata["ss_sd_model_name"] = sd_model_name
|
|
||||||
|
|
||||||
if args.vae is not None:
|
|
||||||
vae_name = args.vae
|
|
||||||
if os.path.exists(vae_name):
|
|
||||||
metadata["ss_vae_hash"] = train_util.model_hash(vae_name)
|
|
||||||
metadata["ss_new_vae_hash"] = train_util.calculate_sha256(vae_name)
|
|
||||||
vae_name = os.path.basename(vae_name)
|
|
||||||
metadata["ss_vae_name"] = vae_name
|
|
||||||
|
|
||||||
metadata = {k: str(v) for k, v in metadata.items()}
|
|
||||||
|
|
||||||
progress_bar = tqdm(range(args.max_train_steps), smoothing=0, disable=not accelerator.is_local_main_process, desc="steps")
|
|
||||||
global_step = 0
|
|
||||||
|
|
||||||
noise_scheduler = DDPMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear",
|
|
||||||
num_train_timesteps=1000, clip_sample=False)
|
|
||||||
|
|
||||||
if accelerator.is_main_process:
|
|
||||||
accelerator.init_trackers("network_train")
|
|
||||||
|
|
||||||
for epoch in range(num_train_epochs):
|
|
||||||
print(f"epoch {epoch+1}/{num_train_epochs}")
|
|
||||||
train_dataset.set_current_epoch(epoch + 1)
|
|
||||||
|
|
||||||
metadata["ss_epoch"] = str(epoch+1)
|
|
||||||
|
|
||||||
network.on_epoch_start(text_encoder, unet)
|
|
||||||
|
|
||||||
loss_total = 0
|
|
||||||
for step, batch in enumerate(train_dataloader):
|
|
||||||
with accelerator.accumulate(network):
|
|
||||||
with torch.no_grad():
|
|
||||||
if "latents" in batch and batch["latents"] is not None:
|
|
||||||
latents = batch["latents"].to(accelerator.device)
|
|
||||||
else:
|
|
||||||
# latentに変換
|
|
||||||
latents = vae.encode(batch["images"].to(dtype=weight_dtype)).latent_dist.sample()
|
|
||||||
hint_latents = vae.encode(batch["hint_images"].to(dtype=weight_dtype)).latent_dist.sample()
|
|
||||||
latents = latents * 0.18215
|
|
||||||
hint_latents = hint_latents * 0.18215
|
|
||||||
# hint = torch.nn.functional.interpolate(batch["hint_images"], scale_factor=(1/8, 1/8), mode="bilinear")
|
|
||||||
# hint = hint[:, 0].unsqueeze(1) # RGB -> BW
|
|
||||||
b_size = latents.shape[0]
|
|
||||||
|
|
||||||
with torch.set_grad_enabled(train_text_encoder):
|
|
||||||
# Get the text embedding for conditioning
|
|
||||||
input_ids = batch["input_ids"].to(accelerator.device)
|
|
||||||
encoder_hidden_states = train_util.get_hidden_states(args, input_ids, tokenizer, text_encoder, weight_dtype)
|
|
||||||
|
|
||||||
# Sample noise that we'll add to the latents
|
|
||||||
noise = torch.randn_like(latents, device=latents.device)
|
|
||||||
|
|
||||||
# Sample a random timestep for each image
|
|
||||||
timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (b_size,), device=latents.device)
|
|
||||||
timesteps = timesteps.long()
|
|
||||||
|
|
||||||
# Add noise to the latents according to the noise magnitude at each timestep
|
|
||||||
# (this is the forward diffusion process)
|
|
||||||
noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
|
|
||||||
|
|
||||||
# Predict the noise residual
|
|
||||||
network.set_as_control_path(True)
|
|
||||||
unet(hint_latents, timesteps, encoder_hidden_states) # めちゃくちゃ乱暴だが入力にhintを加える
|
|
||||||
# unet(noisy_latents * hint, timesteps, encoder_hidden_states) # めちゃくちゃ乱暴だが入力にhintを乗算
|
|
||||||
network.set_as_control_path(False)
|
|
||||||
noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
|
|
||||||
|
|
||||||
if args.v_parameterization:
|
|
||||||
# v-parameterization training
|
|
||||||
target = noise_scheduler.get_velocity(latents, noise, timesteps)
|
|
||||||
else:
|
|
||||||
target = noise
|
|
||||||
|
|
||||||
loss = torch.nn.functional.mse_loss(noise_pred.float(), target.float(), reduction="none")
|
|
||||||
loss = loss.mean([1, 2, 3])
|
|
||||||
|
|
||||||
loss_weights = batch["loss_weights"] # 各sampleごとのweight
|
|
||||||
loss = loss * loss_weights
|
|
||||||
|
|
||||||
loss = loss.mean() # 平均なのでbatch_sizeで割る必要なし
|
|
||||||
|
|
||||||
accelerator.backward(loss)
|
|
||||||
if accelerator.sync_gradients:
|
|
||||||
params_to_clip = network.get_trainable_params()
|
|
||||||
accelerator.clip_grad_norm_(params_to_clip, 1.0) # args.max_grad_norm)
|
|
||||||
|
|
||||||
optimizer.step()
|
|
||||||
lr_scheduler.step()
|
|
||||||
optimizer.zero_grad(set_to_none=True)
|
|
||||||
|
|
||||||
# Checks if the accelerator has performed an optimization step behind the scenes
|
|
||||||
if accelerator.sync_gradients:
|
|
||||||
progress_bar.update(1)
|
|
||||||
global_step += 1
|
|
||||||
|
|
||||||
current_loss = loss.detach().item()
|
|
||||||
loss_total += current_loss
|
|
||||||
avr_loss = loss_total / (step+1)
|
|
||||||
logs = {"loss": avr_loss} # , "lr": lr_scheduler.get_last_lr()[0]}
|
|
||||||
progress_bar.set_postfix(**logs)
|
|
||||||
|
|
||||||
if args.logging_dir is not None:
|
|
||||||
logs = generate_step_logs(args, current_loss, avr_loss, lr_scheduler)
|
|
||||||
accelerator.log(logs, step=global_step)
|
|
||||||
|
|
||||||
if global_step >= args.max_train_steps:
|
|
||||||
break
|
|
||||||
|
|
||||||
if args.logging_dir is not None:
|
|
||||||
logs = {"loss/epoch": loss_total / len(train_dataloader)}
|
|
||||||
accelerator.log(logs, step=epoch+1)
|
|
||||||
|
|
||||||
accelerator.wait_for_everyone()
|
|
||||||
|
|
||||||
if args.save_every_n_epochs is not None:
|
|
||||||
model_name = train_util.DEFAULT_EPOCH_NAME if args.output_name is None else args.output_name
|
|
||||||
|
|
||||||
def save_func():
|
|
||||||
ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, epoch + 1) + '.' + args.save_model_as
|
|
||||||
ckpt_file = os.path.join(args.output_dir, ckpt_name)
|
|
||||||
print(f"saving checkpoint: {ckpt_file}")
|
|
||||||
unwrap_model(network).save_weights(ckpt_file, save_dtype, None if args.no_metadata else metadata)
|
|
||||||
|
|
||||||
def remove_old_func(old_epoch_no):
|
|
||||||
old_ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, old_epoch_no) + '.' + args.save_model_as
|
|
||||||
old_ckpt_file = os.path.join(args.output_dir, old_ckpt_name)
|
|
||||||
if os.path.exists(old_ckpt_file):
|
|
||||||
print(f"removing old checkpoint: {old_ckpt_file}")
|
|
||||||
os.remove(old_ckpt_file)
|
|
||||||
|
|
||||||
saving = train_util.save_on_epoch_end(args, save_func, remove_old_func, epoch + 1, num_train_epochs)
|
|
||||||
if saving and args.save_state:
|
|
||||||
train_util.save_state_on_epoch_end(args, accelerator, model_name, epoch + 1)
|
|
||||||
|
|
||||||
# end of epoch
|
|
||||||
|
|
||||||
metadata["ss_epoch"] = str(num_train_epochs)
|
|
||||||
|
|
||||||
is_main_process = accelerator.is_main_process
|
|
||||||
if is_main_process:
|
|
||||||
network = unwrap_model(network)
|
|
||||||
|
|
||||||
accelerator.end_training()
|
|
||||||
|
|
||||||
if args.save_state:
|
|
||||||
train_util.save_state_on_train_end(args, accelerator)
|
|
||||||
|
|
||||||
del accelerator # この後メモリを使うのでこれは消す
|
|
||||||
|
|
||||||
if is_main_process:
|
|
||||||
os.makedirs(args.output_dir, exist_ok=True)
|
|
||||||
|
|
||||||
model_name = train_util.DEFAULT_LAST_OUTPUT_NAME if args.output_name is None else args.output_name
|
|
||||||
ckpt_name = model_name + '.' + args.save_model_as
|
|
||||||
ckpt_file = os.path.join(args.output_dir, ckpt_name)
|
|
||||||
|
|
||||||
print(f"save trained model to {ckpt_file}")
|
|
||||||
network.save_weights(ckpt_file, save_dtype, None if args.no_metadata else metadata)
|
|
||||||
print("model saved.")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
|
||||||
parser = argparse.ArgumentParser()
|
|
||||||
|
|
||||||
train_util.add_sd_models_arguments(parser)
|
|
||||||
train_util.add_dataset_arguments(parser, True, True, True)
|
|
||||||
train_util.add_training_arguments(parser, True)
|
|
||||||
|
|
||||||
parser.add_argument("--no_metadata", action='store_true', help="do not save metadata in output model / メタデータを出力先モデルに保存しない")
|
|
||||||
parser.add_argument("--save_model_as", type=str, default="safetensors", choices=[None, "ckpt", "pt", "safetensors"],
|
|
||||||
help="format to save the model (default is .safetensors) / モデル保存時の形式（デフォルトはsafetensors）")
|
|
||||||
|
|
||||||
parser.add_argument("--unet_lr", type=float, default=None, help="learning rate for U-Net / U-Netの学習率")
|
|
||||||
parser.add_argument("--text_encoder_lr", type=float, default=None, help="learning rate for Text Encoder / Text Encoderの学習率")
|
|
||||||
parser.add_argument("--lr_scheduler_num_cycles", type=int, default=1,
|
|
||||||
help="Number of restarts for cosine scheduler with restarts / cosine with restartsスケジューラでのリスタート回数")
|
|
||||||
parser.add_argument("--lr_scheduler_power", type=float, default=1,
|
|
||||||
help="Polynomial power for polynomial scheduler / polynomialスケジューラでのpolynomial power")
|
|
||||||
|
|
||||||
parser.add_argument("--network_weights", type=str, default=None,
|
|
||||||
help="pretrained weights for network / 学習するネットワークの初期重み")
|
|
||||||
parser.add_argument("--network_module", type=str, default=None, help='network module to train / 学習対象のネットワークのモジュール')
|
|
||||||
parser.add_argument("--network_dim", type=int, default=None,
|
|
||||||
help='network dimensions (depends on each network) / モジュールの次元数（ネットワークにより定義は異なります）')
|
|
||||||
parser.add_argument("--network_alpha", type=float, default=1,
|
|
||||||
help='alpha for LoRA weight scaling, default 1 (same as network_dim for same behavior as old version) / LoRaの重み調整のalpha値、デフォルト1（旧バージョンと同じ動作をするにはnetwork_dimと同じ値を指定）')
|
|
||||||
parser.add_argument("--network_args", type=str, default=None, nargs='*',
|
|
||||||
help='additional argmuments for network (key=value) / ネットワークへの追加の引数')
|
|
||||||
parser.add_argument("--network_train_unet_only", action="store_true", help="only training U-Net part / U-Net関連部分のみ学習する")
|
|
||||||
parser.add_argument("--network_train_text_encoder_only", action="store_true",
|
|
||||||
help="only training Text Encoder part / Text Encoder関連部分のみ学習する")
|
|
||||||
parser.add_argument("--training_comment", type=str, default=None,
|
|
||||||
help="arbitrary comment string stored in metadata / メタデータに記録する任意のコメント文字列")
|
|
||||||
|
|
||||||
args = parser.parse_args()
|
|
||||||
train(args)
|
|
||||||
Reference in New Issue
Block a user