782 lines
36 KiB
Python
Executable File
782 lines
36 KiB
Python
Executable File
import datetime
|
|
import glob
|
|
import html
|
|
import os
|
|
import inspect
|
|
from contextlib import closing
|
|
|
|
import torch
|
|
import tqdm
|
|
from einops import rearrange, repeat
|
|
from backend.nn.unet import default
|
|
from modules import devices, sd_models, shared, sd_samplers, hashes, sd_hijack_checkpoint, errors
|
|
from modules.textual_inversion import textual_inversion
|
|
from torch import einsum
|
|
from torch.nn.init import normal_, xavier_normal_, xavier_uniform_, kaiming_normal_, kaiming_uniform_, zeros_
|
|
|
|
from collections import deque
|
|
from statistics import stdev, mean
|
|
|
|
|
|
optimizer_dict = {optim_name : cls_obj for optim_name, cls_obj in inspect.getmembers(torch.optim, inspect.isclass) if optim_name != "Optimizer"}
|
|
|
|
class HypernetworkModule(torch.nn.Module):
|
|
activation_dict = {
|
|
"linear": torch.nn.Identity,
|
|
"relu": torch.nn.ReLU,
|
|
"leakyrelu": torch.nn.LeakyReLU,
|
|
"elu": torch.nn.ELU,
|
|
"swish": torch.nn.Hardswish,
|
|
"tanh": torch.nn.Tanh,
|
|
"sigmoid": torch.nn.Sigmoid,
|
|
}
|
|
activation_dict.update({cls_name.lower(): cls_obj for cls_name, cls_obj in inspect.getmembers(torch.nn.modules.activation) if inspect.isclass(cls_obj) and cls_obj.__module__ == 'torch.nn.modules.activation'})
|
|
|
|
def __init__(self, dim, state_dict=None, layer_structure=None, activation_func=None, weight_init='Normal',
|
|
add_layer_norm=False, activate_output=False, dropout_structure=None):
|
|
super().__init__()
|
|
|
|
self.multiplier = 1.0
|
|
|
|
assert layer_structure is not None, "layer_structure must not be None"
|
|
assert layer_structure[0] == 1, "Multiplier Sequence should start with size 1!"
|
|
assert layer_structure[-1] == 1, "Multiplier Sequence should end with size 1!"
|
|
|
|
linears = []
|
|
for i in range(len(layer_structure) - 1):
|
|
|
|
# Add a fully-connected layer
|
|
linears.append(torch.nn.Linear(int(dim * layer_structure[i]), int(dim * layer_structure[i+1])))
|
|
|
|
# Add an activation func except last layer
|
|
if activation_func == "linear" or activation_func is None or (i >= len(layer_structure) - 2 and not activate_output):
|
|
pass
|
|
elif activation_func in self.activation_dict:
|
|
linears.append(self.activation_dict[activation_func]())
|
|
else:
|
|
raise RuntimeError(f'hypernetwork uses an unsupported activation function: {activation_func}')
|
|
|
|
# Add layer normalization
|
|
if add_layer_norm:
|
|
linears.append(torch.nn.LayerNorm(int(dim * layer_structure[i+1])))
|
|
|
|
# Everything should be now parsed into dropout structure, and applied here.
|
|
# Since we only have dropouts after layers, dropout structure should start with 0 and end with 0.
|
|
if dropout_structure is not None and dropout_structure[i+1] > 0:
|
|
assert 0 < dropout_structure[i+1] < 1, "Dropout probability should be 0 or float between 0 and 1!"
|
|
linears.append(torch.nn.Dropout(p=dropout_structure[i+1]))
|
|
# Code explanation : [1, 2, 1] -> dropout is missing when last_layer_dropout is false. [1, 2, 2, 1] -> [0, 0.3, 0, 0], when its True, [0, 0.3, 0.3, 0].
|
|
|
|
self.linear = torch.nn.Sequential(*linears)
|
|
|
|
if state_dict is not None:
|
|
self.fix_old_state_dict(state_dict)
|
|
self.load_state_dict(state_dict)
|
|
else:
|
|
for layer in self.linear:
|
|
if type(layer) == torch.nn.Linear or type(layer) == torch.nn.LayerNorm:
|
|
w, b = layer.weight.data, layer.bias.data
|
|
if weight_init == "Normal" or type(layer) == torch.nn.LayerNorm:
|
|
normal_(w, mean=0.0, std=0.01)
|
|
normal_(b, mean=0.0, std=0)
|
|
elif weight_init == 'XavierUniform':
|
|
xavier_uniform_(w)
|
|
zeros_(b)
|
|
elif weight_init == 'XavierNormal':
|
|
xavier_normal_(w)
|
|
zeros_(b)
|
|
elif weight_init == 'KaimingUniform':
|
|
kaiming_uniform_(w, nonlinearity='leaky_relu' if 'leakyrelu' == activation_func else 'relu')
|
|
zeros_(b)
|
|
elif weight_init == 'KaimingNormal':
|
|
kaiming_normal_(w, nonlinearity='leaky_relu' if 'leakyrelu' == activation_func else 'relu')
|
|
zeros_(b)
|
|
else:
|
|
raise KeyError(f"Key {weight_init} is not defined as initialization!")
|
|
devices.torch_npu_set_device()
|
|
self.to(devices.device)
|
|
|
|
def fix_old_state_dict(self, state_dict):
|
|
changes = {
|
|
'linear1.bias': 'linear.0.bias',
|
|
'linear1.weight': 'linear.0.weight',
|
|
'linear2.bias': 'linear.1.bias',
|
|
'linear2.weight': 'linear.1.weight',
|
|
}
|
|
|
|
for fr, to in changes.items():
|
|
x = state_dict.get(fr, None)
|
|
if x is None:
|
|
continue
|
|
|
|
del state_dict[fr]
|
|
state_dict[to] = x
|
|
|
|
def forward(self, x):
|
|
return x + self.linear(x) * (self.multiplier if not self.training else 1)
|
|
|
|
def trainables(self):
|
|
layer_structure = []
|
|
for layer in self.linear:
|
|
if type(layer) == torch.nn.Linear or type(layer) == torch.nn.LayerNorm:
|
|
layer_structure += [layer.weight, layer.bias]
|
|
return layer_structure
|
|
|
|
|
|
#param layer_structure : sequence used for length, use_dropout : controlling boolean, last_layer_dropout : for compatibility check.
|
|
def parse_dropout_structure(layer_structure, use_dropout, last_layer_dropout):
|
|
if layer_structure is None:
|
|
layer_structure = [1, 2, 1]
|
|
if not use_dropout:
|
|
return [0] * len(layer_structure)
|
|
dropout_values = [0]
|
|
dropout_values.extend([0.3] * (len(layer_structure) - 3))
|
|
if last_layer_dropout:
|
|
dropout_values.append(0.3)
|
|
else:
|
|
dropout_values.append(0)
|
|
dropout_values.append(0)
|
|
return dropout_values
|
|
|
|
|
|
class Hypernetwork:
|
|
filename = None
|
|
name = None
|
|
|
|
def __init__(self, name=None, enable_sizes=None, layer_structure=None, activation_func=None, weight_init=None, add_layer_norm=False, use_dropout=False, activate_output=False, **kwargs):
|
|
self.filename = None
|
|
self.name = name
|
|
self.layers = {}
|
|
self.step = 0
|
|
self.sd_checkpoint = None
|
|
self.sd_checkpoint_name = None
|
|
self.layer_structure = layer_structure
|
|
self.activation_func = activation_func
|
|
self.weight_init = weight_init
|
|
self.add_layer_norm = add_layer_norm
|
|
self.use_dropout = use_dropout
|
|
self.activate_output = activate_output
|
|
self.last_layer_dropout = kwargs.get('last_layer_dropout', True)
|
|
self.dropout_structure = kwargs.get('dropout_structure', None)
|
|
if self.dropout_structure is None:
|
|
self.dropout_structure = parse_dropout_structure(self.layer_structure, self.use_dropout, self.last_layer_dropout)
|
|
self.optimizer_name = None
|
|
self.optimizer_state_dict = None
|
|
self.optional_info = None
|
|
|
|
for size in enable_sizes or []:
|
|
self.layers[size] = (
|
|
HypernetworkModule(size, None, self.layer_structure, self.activation_func, self.weight_init,
|
|
self.add_layer_norm, self.activate_output, dropout_structure=self.dropout_structure),
|
|
HypernetworkModule(size, None, self.layer_structure, self.activation_func, self.weight_init,
|
|
self.add_layer_norm, self.activate_output, dropout_structure=self.dropout_structure),
|
|
)
|
|
self.eval()
|
|
|
|
def weights(self):
|
|
res = []
|
|
for layers in self.layers.values():
|
|
for layer in layers:
|
|
res += layer.parameters()
|
|
return res
|
|
|
|
def train(self, mode=True):
|
|
for layers in self.layers.values():
|
|
for layer in layers:
|
|
layer.train(mode=mode)
|
|
for param in layer.parameters():
|
|
param.requires_grad = mode
|
|
|
|
def to(self, device):
|
|
for layers in self.layers.values():
|
|
for layer in layers:
|
|
layer.to(device)
|
|
|
|
return self
|
|
|
|
def set_multiplier(self, multiplier):
|
|
for layers in self.layers.values():
|
|
for layer in layers:
|
|
layer.multiplier = multiplier
|
|
|
|
return self
|
|
|
|
def eval(self):
|
|
for layers in self.layers.values():
|
|
for layer in layers:
|
|
layer.eval()
|
|
for param in layer.parameters():
|
|
param.requires_grad = False
|
|
|
|
def save(self, filename):
|
|
state_dict = {}
|
|
optimizer_saved_dict = {}
|
|
|
|
for k, v in self.layers.items():
|
|
state_dict[k] = (v[0].state_dict(), v[1].state_dict())
|
|
|
|
state_dict['step'] = self.step
|
|
state_dict['name'] = self.name
|
|
state_dict['layer_structure'] = self.layer_structure
|
|
state_dict['activation_func'] = self.activation_func
|
|
state_dict['is_layer_norm'] = self.add_layer_norm
|
|
state_dict['weight_initialization'] = self.weight_init
|
|
state_dict['sd_checkpoint'] = self.sd_checkpoint
|
|
state_dict['sd_checkpoint_name'] = self.sd_checkpoint_name
|
|
state_dict['activate_output'] = self.activate_output
|
|
state_dict['use_dropout'] = self.use_dropout
|
|
state_dict['dropout_structure'] = self.dropout_structure
|
|
state_dict['last_layer_dropout'] = (self.dropout_structure[-2] != 0) if self.dropout_structure is not None else self.last_layer_dropout
|
|
state_dict['optional_info'] = self.optional_info if self.optional_info else None
|
|
|
|
if self.optimizer_name is not None:
|
|
optimizer_saved_dict['optimizer_name'] = self.optimizer_name
|
|
|
|
torch.save(state_dict, filename)
|
|
if shared.opts.save_optimizer_state and self.optimizer_state_dict:
|
|
optimizer_saved_dict['hash'] = self.shorthash()
|
|
optimizer_saved_dict['optimizer_state_dict'] = self.optimizer_state_dict
|
|
torch.save(optimizer_saved_dict, filename + '.optim')
|
|
|
|
def load(self, filename):
|
|
self.filename = filename
|
|
if self.name is None:
|
|
self.name = os.path.splitext(os.path.basename(filename))[0]
|
|
|
|
state_dict = torch.load(filename, map_location='cpu')
|
|
|
|
self.layer_structure = state_dict.get('layer_structure', [1, 2, 1])
|
|
self.optional_info = state_dict.get('optional_info', None)
|
|
self.activation_func = state_dict.get('activation_func', None)
|
|
self.weight_init = state_dict.get('weight_initialization', 'Normal')
|
|
self.add_layer_norm = state_dict.get('is_layer_norm', False)
|
|
self.dropout_structure = state_dict.get('dropout_structure', None)
|
|
self.use_dropout = True if self.dropout_structure is not None and any(self.dropout_structure) else state_dict.get('use_dropout', False)
|
|
self.activate_output = state_dict.get('activate_output', True)
|
|
self.last_layer_dropout = state_dict.get('last_layer_dropout', False)
|
|
# Dropout structure should have same length as layer structure, Every digits should be in [0,1), and last digit must be 0.
|
|
if self.dropout_structure is None:
|
|
self.dropout_structure = parse_dropout_structure(self.layer_structure, self.use_dropout, self.last_layer_dropout)
|
|
|
|
if shared.opts.print_hypernet_extra:
|
|
if self.optional_info is not None:
|
|
print(f" INFO:\n {self.optional_info}\n")
|
|
|
|
print(f" Layer structure: {self.layer_structure}")
|
|
print(f" Activation function: {self.activation_func}")
|
|
print(f" Weight initialization: {self.weight_init}")
|
|
print(f" Layer norm: {self.add_layer_norm}")
|
|
print(f" Dropout usage: {self.use_dropout}" )
|
|
print(f" Activate last layer: {self.activate_output}")
|
|
print(f" Dropout structure: {self.dropout_structure}")
|
|
|
|
optimizer_saved_dict = torch.load(self.filename + '.optim', map_location='cpu') if os.path.exists(self.filename + '.optim') else {}
|
|
|
|
if self.shorthash() == optimizer_saved_dict.get('hash', None):
|
|
self.optimizer_state_dict = optimizer_saved_dict.get('optimizer_state_dict', None)
|
|
else:
|
|
self.optimizer_state_dict = None
|
|
if self.optimizer_state_dict:
|
|
self.optimizer_name = optimizer_saved_dict.get('optimizer_name', 'AdamW')
|
|
if shared.opts.print_hypernet_extra:
|
|
print("Loaded existing optimizer from checkpoint")
|
|
print(f"Optimizer name is {self.optimizer_name}")
|
|
else:
|
|
self.optimizer_name = "AdamW"
|
|
if shared.opts.print_hypernet_extra:
|
|
print("No saved optimizer exists in checkpoint")
|
|
|
|
for size, sd in state_dict.items():
|
|
if type(size) == int:
|
|
self.layers[size] = (
|
|
HypernetworkModule(size, sd[0], self.layer_structure, self.activation_func, self.weight_init,
|
|
self.add_layer_norm, self.activate_output, self.dropout_structure),
|
|
HypernetworkModule(size, sd[1], self.layer_structure, self.activation_func, self.weight_init,
|
|
self.add_layer_norm, self.activate_output, self.dropout_structure),
|
|
)
|
|
|
|
self.name = state_dict.get('name', self.name)
|
|
self.step = state_dict.get('step', 0)
|
|
self.sd_checkpoint = state_dict.get('sd_checkpoint', None)
|
|
self.sd_checkpoint_name = state_dict.get('sd_checkpoint_name', None)
|
|
self.eval()
|
|
|
|
def shorthash(self):
|
|
sha256 = hashes.sha256(self.filename, f'hypernet/{self.name}')
|
|
|
|
return sha256[0:10] if sha256 else None
|
|
|
|
|
|
def list_hypernetworks(path):
|
|
res = {}
|
|
for filename in sorted(glob.iglob(os.path.join(path, '**/*.pt'), recursive=True), key=str.lower):
|
|
name = os.path.splitext(os.path.basename(filename))[0]
|
|
# Prevent a hypothetical "None.pt" from being listed.
|
|
if name != "None":
|
|
res[name] = filename
|
|
return res
|
|
|
|
|
|
def load_hypernetwork(name):
|
|
path = shared.hypernetworks.get(name, None)
|
|
|
|
if path is None:
|
|
return None
|
|
|
|
try:
|
|
hypernetwork = Hypernetwork()
|
|
hypernetwork.load(path)
|
|
return hypernetwork
|
|
except Exception:
|
|
errors.report(f"Error loading hypernetwork {path}", exc_info=True)
|
|
return None
|
|
|
|
|
|
def load_hypernetworks(names, multipliers=None):
|
|
already_loaded = {}
|
|
|
|
for hypernetwork in shared.loaded_hypernetworks:
|
|
if hypernetwork.name in names:
|
|
already_loaded[hypernetwork.name] = hypernetwork
|
|
|
|
shared.loaded_hypernetworks.clear()
|
|
|
|
for i, name in enumerate(names):
|
|
hypernetwork = already_loaded.get(name, None)
|
|
if hypernetwork is None:
|
|
hypernetwork = load_hypernetwork(name)
|
|
|
|
if hypernetwork is None:
|
|
continue
|
|
|
|
hypernetwork.set_multiplier(multipliers[i] if multipliers else 1.0)
|
|
shared.loaded_hypernetworks.append(hypernetwork)
|
|
|
|
|
|
def apply_single_hypernetwork(hypernetwork, context_k, context_v, layer=None):
|
|
hypernetwork_layers = (hypernetwork.layers if hypernetwork is not None else {}).get(context_k.shape[2], None)
|
|
|
|
if hypernetwork_layers is None:
|
|
return context_k, context_v
|
|
|
|
if layer is not None:
|
|
layer.hyper_k = hypernetwork_layers[0]
|
|
layer.hyper_v = hypernetwork_layers[1]
|
|
|
|
context_k = devices.cond_cast_unet(hypernetwork_layers[0](devices.cond_cast_float(context_k)))
|
|
context_v = devices.cond_cast_unet(hypernetwork_layers[1](devices.cond_cast_float(context_v)))
|
|
return context_k, context_v
|
|
|
|
|
|
def apply_hypernetworks(hypernetworks, context, layer=None):
|
|
context_k = context
|
|
context_v = context
|
|
for hypernetwork in hypernetworks:
|
|
context_k, context_v = apply_single_hypernetwork(hypernetwork, context_k, context_v, layer)
|
|
|
|
return context_k, context_v
|
|
|
|
|
|
def attention_CrossAttention_forward(self, x, context=None, mask=None, **kwargs):
|
|
h = self.heads
|
|
|
|
q = self.to_q(x)
|
|
context = default(context, x)
|
|
|
|
context_k, context_v = apply_hypernetworks(shared.loaded_hypernetworks, context, self)
|
|
k = self.to_k(context_k)
|
|
v = self.to_v(context_v)
|
|
|
|
q, k, v = (rearrange(t, 'b n (h d) -> (b h) n d', h=h) for t in (q, k, v))
|
|
|
|
sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
|
|
|
|
if mask is not None:
|
|
mask = rearrange(mask, 'b ... -> b (...)')
|
|
max_neg_value = -torch.finfo(sim.dtype).max
|
|
mask = repeat(mask, 'b j -> (b h) () j', h=h)
|
|
sim.masked_fill_(~mask, max_neg_value)
|
|
|
|
# attention, what we cannot get enough of
|
|
attn = sim.softmax(dim=-1)
|
|
|
|
out = einsum('b i j, b j d -> b i d', attn, v)
|
|
out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
|
|
return self.to_out(out)
|
|
|
|
|
|
def stack_conds(conds):
|
|
if len(conds) == 1:
|
|
return torch.stack(conds)
|
|
|
|
# same as in reconstruct_multicond_batch
|
|
token_count = max([x.shape[0] for x in conds])
|
|
for i in range(len(conds)):
|
|
if conds[i].shape[0] != token_count:
|
|
last_vector = conds[i][-1:]
|
|
last_vector_repeated = last_vector.repeat([token_count - conds[i].shape[0], 1])
|
|
conds[i] = torch.vstack([conds[i], last_vector_repeated])
|
|
|
|
return torch.stack(conds)
|
|
|
|
|
|
def statistics(data):
|
|
if len(data) < 2:
|
|
std = 0
|
|
else:
|
|
std = stdev(data)
|
|
total_information = f"loss:{mean(data):.3f}" + u"\u00B1" + f"({std/ (len(data) ** 0.5):.3f})"
|
|
recent_data = data[-32:]
|
|
if len(recent_data) < 2:
|
|
std = 0
|
|
else:
|
|
std = stdev(recent_data)
|
|
recent_information = f"recent 32 loss:{mean(recent_data):.3f}" + u"\u00B1" + f"({std / (len(recent_data) ** 0.5):.3f})"
|
|
return total_information, recent_information
|
|
|
|
#
|
|
# def create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure=None, activation_func=None, weight_init=None, add_layer_norm=False, use_dropout=False, dropout_structure=None):
|
|
# # Remove illegal characters from name.
|
|
# name = "".join( x for x in name if (x.isalnum() or x in "._- "))
|
|
# assert name, "Name cannot be empty!"
|
|
#
|
|
# fn = os.path.join(shared.cmd_opts.hypernetwork_dir, f"{name}.pt")
|
|
# if not overwrite_old:
|
|
# assert not os.path.exists(fn), f"file {fn} already exists"
|
|
#
|
|
# if type(layer_structure) == str:
|
|
# layer_structure = [float(x.strip()) for x in layer_structure.split(",")]
|
|
#
|
|
# if use_dropout and dropout_structure and type(dropout_structure) == str:
|
|
# dropout_structure = [float(x.strip()) for x in dropout_structure.split(",")]
|
|
# else:
|
|
# dropout_structure = [0] * len(layer_structure)
|
|
#
|
|
# hypernet = modules.hypernetworks.hypernetwork.Hypernetwork(
|
|
# name=name,
|
|
# enable_sizes=[int(x) for x in enable_sizes],
|
|
# layer_structure=layer_structure,
|
|
# activation_func=activation_func,
|
|
# weight_init=weight_init,
|
|
# add_layer_norm=add_layer_norm,
|
|
# use_dropout=use_dropout,
|
|
# dropout_structure=dropout_structure
|
|
# )
|
|
# hypernet.save(fn)
|
|
#
|
|
# shared.reload_hypernetworks()
|
|
#
|
|
#
|
|
# def train_hypernetwork(id_task, hypernetwork_name: str, learn_rate: float, batch_size: int, gradient_step: int, data_root: str, log_directory: str, training_width: int, training_height: int, varsize: bool, steps: int, clip_grad_mode: str, clip_grad_value: float, shuffle_tags: bool, tag_drop_out: bool, latent_sampling_method: str, use_weight: bool, create_image_every: int, save_hypernetwork_every: int, template_filename: str, preview_from_txt2img: bool, preview_prompt: str, preview_negative_prompt: str, preview_steps: int, preview_sampler_name: str, preview_cfg_scale: float, preview_seed: int, preview_width: int, preview_height: int):
|
|
# from modules import images, processing
|
|
#
|
|
# save_hypernetwork_every = save_hypernetwork_every or 0
|
|
# create_image_every = create_image_every or 0
|
|
# template_file = textual_inversion.textual_inversion_templates.get(template_filename, None)
|
|
# textual_inversion.validate_train_inputs(hypernetwork_name, learn_rate, batch_size, gradient_step, data_root, template_file, template_filename, steps, save_hypernetwork_every, create_image_every, log_directory, name="hypernetwork")
|
|
# template_file = template_file.path
|
|
#
|
|
# path = shared.hypernetworks.get(hypernetwork_name, None)
|
|
# hypernetwork = Hypernetwork()
|
|
# hypernetwork.load(path)
|
|
# shared.loaded_hypernetworks = [hypernetwork]
|
|
#
|
|
# shared.state.job = "train-hypernetwork"
|
|
# shared.state.textinfo = "Initializing hypernetwork training..."
|
|
# shared.state.job_count = steps
|
|
#
|
|
# hypernetwork_name = hypernetwork_name.rsplit('(', 1)[0]
|
|
# filename = os.path.join(shared.cmd_opts.hypernetwork_dir, f'{hypernetwork_name}.pt')
|
|
#
|
|
# log_directory = os.path.join(log_directory, datetime.datetime.now().strftime("%Y-%m-%d"), hypernetwork_name)
|
|
# unload = shared.opts.unload_models_when_training
|
|
#
|
|
# if save_hypernetwork_every > 0:
|
|
# hypernetwork_dir = os.path.join(log_directory, "hypernetworks")
|
|
# os.makedirs(hypernetwork_dir, exist_ok=True)
|
|
# else:
|
|
# hypernetwork_dir = None
|
|
#
|
|
# if create_image_every > 0:
|
|
# images_dir = os.path.join(log_directory, "images")
|
|
# os.makedirs(images_dir, exist_ok=True)
|
|
# else:
|
|
# images_dir = None
|
|
#
|
|
# checkpoint = sd_models.select_checkpoint()
|
|
#
|
|
# initial_step = hypernetwork.step or 0
|
|
# if initial_step >= steps:
|
|
# shared.state.textinfo = "Model has already been trained beyond specified max steps"
|
|
# return hypernetwork, filename
|
|
#
|
|
# scheduler = LearnRateScheduler(learn_rate, steps, initial_step)
|
|
#
|
|
# clip_grad = torch.nn.utils.clip_grad_value_ if clip_grad_mode == "value" else torch.nn.utils.clip_grad_norm_ if clip_grad_mode == "norm" else None
|
|
# if clip_grad:
|
|
# clip_grad_sched = LearnRateScheduler(clip_grad_value, steps, initial_step, verbose=False)
|
|
#
|
|
# if shared.opts.training_enable_tensorboard:
|
|
# tensorboard_writer = textual_inversion.tensorboard_setup(log_directory)
|
|
#
|
|
# # dataset loading may take a while, so input validations and early returns should be done before this
|
|
# shared.state.textinfo = f"Preparing dataset from {html.escape(data_root)}..."
|
|
#
|
|
# pin_memory = shared.opts.pin_memory
|
|
#
|
|
# ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=hypernetwork_name, model=shared.sd_model, cond_model=shared.sd_model.cond_stage_model, device=devices.device, template_file=template_file, include_cond=True, batch_size=batch_size, gradient_step=gradient_step, shuffle_tags=shuffle_tags, tag_drop_out=tag_drop_out, latent_sampling_method=latent_sampling_method, varsize=varsize, use_weight=use_weight)
|
|
#
|
|
# if shared.opts.save_training_settings_to_txt:
|
|
# saved_params = dict(
|
|
# model_name=checkpoint.model_name, model_hash=checkpoint.shorthash, num_of_dataset_images=len(ds),
|
|
# **{field: getattr(hypernetwork, field) for field in ['layer_structure', 'activation_func', 'weight_init', 'add_layer_norm', 'use_dropout', ]}
|
|
# )
|
|
# saving_settings.save_settings_to_file(log_directory, {**saved_params, **locals()})
|
|
#
|
|
# latent_sampling_method = ds.latent_sampling_method
|
|
#
|
|
# dl = modules.textual_inversion.dataset.PersonalizedDataLoader(ds, latent_sampling_method=latent_sampling_method, batch_size=ds.batch_size, pin_memory=pin_memory)
|
|
#
|
|
# old_parallel_processing_allowed = shared.parallel_processing_allowed
|
|
#
|
|
# if unload:
|
|
# shared.parallel_processing_allowed = False
|
|
# shared.sd_model.cond_stage_model.to(devices.cpu)
|
|
# shared.sd_model.first_stage_model.to(devices.cpu)
|
|
#
|
|
# weights = hypernetwork.weights()
|
|
# hypernetwork.train()
|
|
#
|
|
# # Here we use optimizer from saved HN, or we can specify as UI option.
|
|
# if hypernetwork.optimizer_name in optimizer_dict:
|
|
# optimizer = optimizer_dict[hypernetwork.optimizer_name](params=weights, lr=scheduler.learn_rate)
|
|
# optimizer_name = hypernetwork.optimizer_name
|
|
# else:
|
|
# print(f"Optimizer type {hypernetwork.optimizer_name} is not defined!")
|
|
# optimizer = torch.optim.AdamW(params=weights, lr=scheduler.learn_rate)
|
|
# optimizer_name = 'AdamW'
|
|
#
|
|
# if hypernetwork.optimizer_state_dict: # This line must be changed if Optimizer type can be different from saved optimizer.
|
|
# try:
|
|
# optimizer.load_state_dict(hypernetwork.optimizer_state_dict)
|
|
# except RuntimeError as e:
|
|
# print("Cannot resume from saved optimizer!")
|
|
# print(e)
|
|
#
|
|
# scaler = torch.cuda.amp.GradScaler()
|
|
#
|
|
# batch_size = ds.batch_size
|
|
# gradient_step = ds.gradient_step
|
|
# # n steps = batch_size * gradient_step * n image processed
|
|
# steps_per_epoch = len(ds) // batch_size // gradient_step
|
|
# max_steps_per_epoch = len(ds) // batch_size - (len(ds) // batch_size) % gradient_step
|
|
# loss_step = 0
|
|
# _loss_step = 0 #internal
|
|
# # size = len(ds.indexes)
|
|
# # loss_dict = defaultdict(lambda : deque(maxlen = 1024))
|
|
# loss_logging = deque(maxlen=len(ds) * 3) # this should be configurable parameter, this is 3 * epoch(dataset size)
|
|
# # losses = torch.zeros((size,))
|
|
# # previous_mean_losses = [0]
|
|
# # previous_mean_loss = 0
|
|
# # print("Mean loss of {} elements".format(size))
|
|
#
|
|
# steps_without_grad = 0
|
|
#
|
|
# last_saved_file = "<none>"
|
|
# last_saved_image = "<none>"
|
|
# forced_filename = "<none>"
|
|
#
|
|
# pbar = tqdm.tqdm(total=steps - initial_step)
|
|
# try:
|
|
# sd_hijack_checkpoint.add()
|
|
#
|
|
# for _ in range((steps-initial_step) * gradient_step):
|
|
# if scheduler.finished:
|
|
# break
|
|
# if shared.state.interrupted:
|
|
# break
|
|
# for j, batch in enumerate(dl):
|
|
# # works as a drop_last=True for gradient accumulation
|
|
# if j == max_steps_per_epoch:
|
|
# break
|
|
# scheduler.apply(optimizer, hypernetwork.step)
|
|
# if scheduler.finished:
|
|
# break
|
|
# if shared.state.interrupted:
|
|
# break
|
|
#
|
|
# if clip_grad:
|
|
# clip_grad_sched.step(hypernetwork.step)
|
|
#
|
|
# with devices.autocast():
|
|
# x = batch.latent_sample.to(devices.device, non_blocking=pin_memory)
|
|
# if use_weight:
|
|
# w = batch.weight.to(devices.device, non_blocking=pin_memory)
|
|
# if tag_drop_out != 0 or shuffle_tags:
|
|
# shared.sd_model.cond_stage_model.to(devices.device)
|
|
# c = shared.sd_model.cond_stage_model(batch.cond_text).to(devices.device, non_blocking=pin_memory)
|
|
# shared.sd_model.cond_stage_model.to(devices.cpu)
|
|
# else:
|
|
# c = stack_conds(batch.cond).to(devices.device, non_blocking=pin_memory)
|
|
# if use_weight:
|
|
# loss = shared.sd_model.weighted_forward(x, c, w)[0] / gradient_step
|
|
# del w
|
|
# else:
|
|
# loss = shared.sd_model.forward(x, c)[0] / gradient_step
|
|
# del x
|
|
# del c
|
|
#
|
|
# _loss_step += loss.item()
|
|
# scaler.scale(loss).backward()
|
|
#
|
|
# # go back until we reach gradient accumulation steps
|
|
# if (j + 1) % gradient_step != 0:
|
|
# continue
|
|
# loss_logging.append(_loss_step)
|
|
# if clip_grad:
|
|
# clip_grad(weights, clip_grad_sched.learn_rate)
|
|
#
|
|
# scaler.step(optimizer)
|
|
# scaler.update()
|
|
# hypernetwork.step += 1
|
|
# pbar.update()
|
|
# optimizer.zero_grad(set_to_none=True)
|
|
# loss_step = _loss_step
|
|
# _loss_step = 0
|
|
#
|
|
# steps_done = hypernetwork.step + 1
|
|
#
|
|
# epoch_num = hypernetwork.step // steps_per_epoch
|
|
# epoch_step = hypernetwork.step % steps_per_epoch
|
|
#
|
|
# description = f"Training hypernetwork [Epoch {epoch_num}: {epoch_step+1}/{steps_per_epoch}]loss: {loss_step:.7f}"
|
|
# pbar.set_description(description)
|
|
# if hypernetwork_dir is not None and steps_done % save_hypernetwork_every == 0:
|
|
# # Before saving, change name to match current checkpoint.
|
|
# hypernetwork_name_every = f'{hypernetwork_name}-{steps_done}'
|
|
# last_saved_file = os.path.join(hypernetwork_dir, f'{hypernetwork_name_every}.pt')
|
|
# hypernetwork.optimizer_name = optimizer_name
|
|
# if shared.opts.save_optimizer_state:
|
|
# hypernetwork.optimizer_state_dict = optimizer.state_dict()
|
|
# save_hypernetwork(hypernetwork, checkpoint, hypernetwork_name, last_saved_file)
|
|
# hypernetwork.optimizer_state_dict = None # dereference it after saving, to save memory.
|
|
#
|
|
#
|
|
#
|
|
# if shared.opts.training_enable_tensorboard:
|
|
# epoch_num = hypernetwork.step // len(ds)
|
|
# epoch_step = hypernetwork.step - (epoch_num * len(ds)) + 1
|
|
# mean_loss = sum(loss_logging) / len(loss_logging)
|
|
# textual_inversion.tensorboard_add(tensorboard_writer, loss=mean_loss, global_step=hypernetwork.step, step=epoch_step, learn_rate=scheduler.learn_rate, epoch_num=epoch_num)
|
|
#
|
|
# textual_inversion.write_loss(log_directory, "hypernetwork_loss.csv", hypernetwork.step, steps_per_epoch, {
|
|
# "loss": f"{loss_step:.7f}",
|
|
# "learn_rate": scheduler.learn_rate
|
|
# })
|
|
#
|
|
# if images_dir is not None and steps_done % create_image_every == 0:
|
|
# forced_filename = f'{hypernetwork_name}-{steps_done}'
|
|
# last_saved_image = os.path.join(images_dir, forced_filename)
|
|
# hypernetwork.eval()
|
|
# rng_state = torch.get_rng_state()
|
|
# cuda_rng_state = None
|
|
# if torch.cuda.is_available():
|
|
# cuda_rng_state = torch.cuda.get_rng_state_all()
|
|
# shared.sd_model.cond_stage_model.to(devices.device)
|
|
# shared.sd_model.first_stage_model.to(devices.device)
|
|
#
|
|
# p = processing.StableDiffusionProcessingTxt2Img(
|
|
# sd_model=shared.sd_model,
|
|
# do_not_save_grid=True,
|
|
# do_not_save_samples=True,
|
|
# )
|
|
#
|
|
# p.disable_extra_networks = True
|
|
#
|
|
# if preview_from_txt2img:
|
|
# p.prompt = preview_prompt
|
|
# p.negative_prompt = preview_negative_prompt
|
|
# p.steps = preview_steps
|
|
# p.sampler_name = sd_samplers.samplers_map[preview_sampler_name.lower()]
|
|
# p.cfg_scale = preview_cfg_scale
|
|
# p.seed = preview_seed
|
|
# p.width = preview_width
|
|
# p.height = preview_height
|
|
# else:
|
|
# p.prompt = batch.cond_text[0]
|
|
# p.steps = 20
|
|
# p.width = training_width
|
|
# p.height = training_height
|
|
#
|
|
# preview_text = p.prompt
|
|
#
|
|
# with closing(p):
|
|
# processed = processing.process_images(p)
|
|
# image = processed.images[0] if len(processed.images) > 0 else None
|
|
#
|
|
# if unload:
|
|
# shared.sd_model.cond_stage_model.to(devices.cpu)
|
|
# shared.sd_model.first_stage_model.to(devices.cpu)
|
|
# torch.set_rng_state(rng_state)
|
|
# if torch.cuda.is_available():
|
|
# torch.cuda.set_rng_state_all(cuda_rng_state)
|
|
# hypernetwork.train()
|
|
# if image is not None:
|
|
# shared.state.assign_current_image(image)
|
|
# if shared.opts.training_enable_tensorboard and shared.opts.training_tensorboard_save_images:
|
|
# textual_inversion.tensorboard_add_image(tensorboard_writer,
|
|
# f"Validation at epoch {epoch_num}", image,
|
|
# hypernetwork.step)
|
|
# last_saved_image, last_text_info = images.save_image(image, images_dir, "", p.seed, p.prompt, shared.opts.samples_format, processed.infotexts[0], p=p, forced_filename=forced_filename, save_to_dirs=False)
|
|
# last_saved_image += f", prompt: {preview_text}"
|
|
#
|
|
# shared.state.job_no = hypernetwork.step
|
|
#
|
|
# shared.state.textinfo = f"""
|
|
# <p>
|
|
# Loss: {loss_step:.7f}<br/>
|
|
# Step: {steps_done}<br/>
|
|
# Last prompt: {html.escape(batch.cond_text[0])}<br/>
|
|
# Last saved hypernetwork: {html.escape(last_saved_file)}<br/>
|
|
# Last saved image: {html.escape(last_saved_image)}<br/>
|
|
# </p>
|
|
# """
|
|
# except Exception:
|
|
# errors.report("Exception in training hypernetwork", exc_info=True)
|
|
# finally:
|
|
# pbar.leave = False
|
|
# pbar.close()
|
|
# hypernetwork.eval()
|
|
# sd_hijack_checkpoint.remove()
|
|
#
|
|
#
|
|
#
|
|
# filename = os.path.join(shared.cmd_opts.hypernetwork_dir, f'{hypernetwork_name}.pt')
|
|
# hypernetwork.optimizer_name = optimizer_name
|
|
# if shared.opts.save_optimizer_state:
|
|
# hypernetwork.optimizer_state_dict = optimizer.state_dict()
|
|
# save_hypernetwork(hypernetwork, checkpoint, hypernetwork_name, filename)
|
|
#
|
|
# del optimizer
|
|
# hypernetwork.optimizer_state_dict = None # dereference it after saving, to save memory.
|
|
# shared.sd_model.cond_stage_model.to(devices.device)
|
|
# shared.sd_model.first_stage_model.to(devices.device)
|
|
# shared.parallel_processing_allowed = old_parallel_processing_allowed
|
|
#
|
|
# return hypernetwork, filename
|
|
#
|
|
# def save_hypernetwork(hypernetwork, checkpoint, hypernetwork_name, filename):
|
|
# old_hypernetwork_name = hypernetwork.name
|
|
# old_sd_checkpoint = hypernetwork.sd_checkpoint if hasattr(hypernetwork, "sd_checkpoint") else None
|
|
# old_sd_checkpoint_name = hypernetwork.sd_checkpoint_name if hasattr(hypernetwork, "sd_checkpoint_name") else None
|
|
# try:
|
|
# hypernetwork.sd_checkpoint = checkpoint.shorthash
|
|
# hypernetwork.sd_checkpoint_name = checkpoint.model_name
|
|
# hypernetwork.name = hypernetwork_name
|
|
# hypernetwork.save(filename)
|
|
# except:
|
|
# hypernetwork.sd_checkpoint = old_sd_checkpoint
|
|
# hypernetwork.sd_checkpoint_name = old_sd_checkpoint_name
|
|
# hypernetwork.name = old_hypernetwork_name
|
|
# raise
|