mirror of
https://github.com/lllyasviel/stable-diffusion-webui-forge.git
synced 2026-03-13 08:59:51 +00:00
remove everything related to training
This commit is contained in:
layerdiffusion
@@ -5,14 +5,12 @@ import os
|
||||
import inspect
|
||||
from contextlib import closing
|
||||
|
||||
import modules.textual_inversion.dataset
|
||||
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, saving_settings
|
||||
from modules.textual_inversion.learn_schedule import LearnRateScheduler
|
||||
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_
|
||||
|
||||
@@ -436,348 +434,348 @@ def statistics(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
|
||||
#
|
||||
# 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
|
||||
|
||||
Reference in New Issue
Block a user