def __init__(self):

self.vae_scale_factor = 0.18215

self.is_sdxl = False

self.interrupted = False

# TODO 他のスクリプトと共通化する

def generate_step_logs(

self,

args: argparse.Namespace,

current_loss,

avr_loss,

lr_scheduler,

lr_descriptions,

keys_scaled=None,

mean_norm=None,

maximum_norm=None,

):

logs = {"loss/current": current_loss, "loss/average": avr_loss}

if keys_scaled is not None:

logs["max_norm/keys_scaled"] = keys_scaled

logs["max_norm/average_key_norm"] = mean_norm

logs["max_norm/max_key_norm"] = maximum_norm

lrs = lr_scheduler.get_last_lr()

for i, lr in enumerate(lrs):

if lr_descriptions is not None:

lr_desc = lr_descriptions[i]

else:

idx = i - (0 if args.network_train_unet_only else -1)

if idx == -1:

lr_desc = "textencoder"

else:

if len(lrs) > 2:

lr_desc = f"group{idx}"

else:

lr_desc = "unet"

logs[f"lr/{lr_desc}"] = lr

if args.optimizer_type.lower().startswith("DAdapt".lower()) or args.optimizer_type.lower() == "Prodigy".lower():

# tracking d*lr value

logs[f"lr/d*lr/{lr_desc}"] = (

lr_scheduler.optimizers[-1].param_groups[i]["d"] * lr_scheduler.optimizers[-1].param_groups[i]["lr"]

)

elif args.mechanize:

logs["lr/s*lr"] = (

torch.sum(lr_scheduler.optimizers[-1].state["_mechanic"]["s"]).item() * lr_scheduler.optimizers[-1].param_groups[0]["lr"]

)

return logs

def assert_extra_args(self, args, train_dataset_group):

pass

def load_target_model(self, args, weight_dtype, accelerator):

text_encoder, vae, unet, _ = train_util.load_target_model(args, weight_dtype, accelerator)

return model_util.get_model_version_str_for_sd1_sd2(args.v2, args.v_parameterization), text_encoder, vae, unet

def load_tokenizer(self, args):

tokenizer = train_util.load_tokenizer(args)

return tokenizer

def is_text_encoder_outputs_cached(self, args):

return False

def is_train_text_encoder(self, args):

return not args.network_train_unet_only and not self.is_text_encoder_outputs_cached(args)

def cache_text_encoder_outputs_if_needed(

self, args, accelerator, unet, vae, tokenizers, text_encoders, data_loader, weight_dtype

):

return

# for t_enc in text_encoders:

# t_enc.to(accelerator.device, dtype=weight_dtype)

def get_text_cond(self, args, accelerator, batch, tokenizers, text_encoders, weight_dtype):

input_ids = batch["input_ids"].to(accelerator.device)

encoder_hidden_states = train_util.get_hidden_states(args, input_ids, tokenizers[0], text_encoders[0], weight_dtype)

return encoder_hidden_states

def call_unet(self, args, accelerator, unet, noisy_latents, timesteps, text_conds, batch, weight_dtype):

noise_pred = unet(noisy_latents, timesteps, text_conds).sample

return noise_pred

def all_reduce_network(self, accelerator, network):

for param in network.parameters():

if param.grad is not None:

param.grad = accelerator.reduce(param.grad, reduction="mean")

def sample_images(self, accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet):

if args.todo_factor:

token_downsampling.remove_patch(unet)

train_util.sample_images(accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet)

if args.todo_factor:

token_downsampling.apply_patch(unet, args, self.is_sdxl)

def create_embedding_from_data(self, data, name='unknown'):

if 'string_to_param' in data: # textual inversion embeddings

param_dict = data['string_to_param']

param_dict = getattr(param_dict, '_parameters', param_dict) # fix for torch 1.12.1 loading saved file from torch 1.11

assert (

len(param_dict) == 1

), f"embedding file has multiple terms in it: {name}"

emb = next(iter(param_dict.items()))[1]

vec = emb.detach()

shape = vec.shape[-1]

vectors = vec.shape[0]

elif type(data) == dict and 'clip_g' in data and 'clip_l' in data: # SDXL embedding

vec = {k: v.detach() for k, v in data.items()}

shape = data['clip_g'].shape[-1] + data['clip_l'].shape[-1]

vectors = data['clip_g'].shape[0]

elif type(data) == dict and type(next(iter(data.values()))) == torch.Tensor: # diffuser concepts

assert (

len(data.keys()) == 1

), f"embedding file has multiple terms in it: {name}"

emb = next(iter(data.values()))

if len(emb.shape) == 1:

emb = emb.unsqueeze(0)

vec = emb.detach()

shape = vec.shape[-1]

vectors = vec.shape[0]

else:

raise Exception(f"Couldn't identify {name} as neither textual inversion embedding nor diffuser concept.")

return vec, shape, vectors

def train(self, 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)

deepspeed_utils.prepare_deepspeed_args(args)

setup_logging(args, reset=True)

cache_latents = args.cache_latents

use_dreambooth_method = args.in_json is None

use_user_config = args.dataset_config is not None

if args.seed is None:

args.seed = random.randint(0, 2**32)

set_seed(args.seed)

# tokenizerは単体またはリスト、tokenizersは必ずリスト：既存のコードとの互換性のため

tokenizer = self.load_tokenizer(args)

tokenizers = tokenizer if isinstance(tokenizer, list) else [tokenizer]

# acceleratorを準備する

logger.info("preparing accelerator")

accelerator = train_util.prepare_accelerator(args)

is_main_process = accelerator.is_main_process

# mixed precisionに対応した型を用意しておき適宜castする

weight_dtype, save_dtype = train_util.prepare_dtype(args)

vae_dtype = torch.float32 if args.no_half_vae else weight_dtype

# モデルを読み込む

apply_low_precision_norm()

model_version, text_encoder, vae, unet = self.load_target_model(args, weight_dtype, accelerator)

undo_low_precision_norm()

# apply token merging patch

if args.todo_factor:

token_downsampling.apply_patch(unet, args, self.is_sdxl)

logger.info(f"enable token downsampling optimization: downsample_factor={args.todo_factor}, max_depth={args.todo_max_depth}")

# text_encoder is List[CLIPTextModel] or CLIPTextModel

text_encoders = text_encoder if isinstance(text_encoder, list) else [text_encoder]

# モデルに xformers とか memory efficient attention を組み込む

train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers, args.sdpa)

if torch.__version__ >= "2.0.0": # PyTorch 2.0.0 以上対応のxformersなら以下が使える

vae.set_use_memory_efficient_attention_xformers(args.xformers)

# prepare embeddings (for pivotal tuning)

# load from existing files

embeddings_map = {}

embedding_to_token_ids = {}

if len(args.embeddings) > 0:

if not args.token_strings:

args.token_strings = [Path(embeds_file).stem for embeds_file in args.embeddings]

if len(args.token_strings) != len(args.embeddings):

raise ValueError(f"token_strings must have a name for each embedding. Got {len(args.token_strings)}/{len(args.embeddings)}")

for embeds_file, token_string in zip(args.embeddings, args.token_strings):

if model_util.is_safetensors(embeds_file):

from safetensors.torch import load_file

data = load_file(embeds_file)

else:

data = torch.load(embeds_file, map_location="cpu")

embeds_list, _shape, num_vectors_per_token = self.create_embedding_from_data(data, token_string)

if isinstance(embeds_list, dict) and "clip_l" in embeds_list and "clip_g" in embeds_list:

embeds_list = [embeds_list["clip_l"], embeds_list["clip_g"]]

else:

embeds_list = [embeds_list]

embedding_to_token_ids[token_string] = []

token_strings = [token_string] + [f"{token_string}{i+1}" for i in range(num_vectors_per_token - 1)]

logger.info(f"Loaded token strings {token_strings}")

for i, (tokenizer, t_enc, embeds) in enumerate(zip(tokenizers, text_encoders, embeds_list)):

num_added_tokens = tokenizer.add_tokens(token_strings)

assert (

num_added_tokens == num_vectors_per_token

), f"The tokenizer already contains {token_string}. Please pass a different word that is not already in the tokenizer. / 指定した名前（ファイル名）のトークンが既に存在します。ファイルをリネームしてください: {embeds_file}"

token_ids = tokenizer.convert_tokens_to_ids(token_strings)

logger.info(f"Textual Inversion embeddings `{token_string}` loaded. Tokens are added: {token_ids}")

assert (

min(token_ids) == token_ids[0] and token_ids[-1] == token_ids[0] + len(token_ids) - 1

), f"token ids is not ordered : tokenizer {i+1}, {token_ids}"

assert (

len(tokenizer) - 1 == token_ids[-1]

), f"token ids is not end of tokenize: tokenizer {i+1}, {token_ids}, {len(tokenizer)}"

# Resize the token embeddings as we are adding new special tokens to the tokenizer

t_enc.resize_token_embeddings(len(tokenizer))

# Initialise the newly added token with the provided embeddings

token_embeds = t_enc.get_input_embeddings().weight.data

for token_id, embed in zip(token_ids, embeds):

token_embeds[token_id] = embed

embedding_to_token_ids[token_string].append(token_ids)

embeddings_map[token_string] = embeds_list

# データセットを準備する

if args.dataset_class is None:

blueprint_generator = BlueprintGenerator(ConfigSanitizer(True, True, args.masked_loss, True))

if use_user_config:

logger.info(f"Loading dataset config from {args.dataset_config}")

user_config = config_util.load_user_config(args.dataset_config)

ignored = ["train_data_dir", "reg_data_dir", "in_json"]

if any(getattr(args, attr) is not None for attr in ignored):

logger.warning(

"ignoring the following options because config file is found: {0} / 設定ファイルが利用されるため以下のオプションは無視されます: {0}".format(

", ".join(ignored)

)

)

else:

if use_dreambooth_method:

logger.info("Using DreamBooth method.")

user_config = {

"datasets": [

{

"subsets": config_util.generate_dreambooth_subsets_config_by_subdirs(

args.train_data_dir, args.reg_data_dir

)

}

]

}

else:

logger.info("Training with captions.")

user_config = {

"datasets": [

{

"subsets": [

{

"image_dir": args.train_data_dir,

"metadata_file": args.in_json,

}

]

}

]

}

blueprint = blueprint_generator.generate(user_config, args, tokenizer=tokenizers)

train_dataset_group = config_util.generate_dataset_group_by_blueprint(blueprint.dataset_group)

else:

# use arbitrary dataset class

train_dataset_group = train_util.load_arbitrary_dataset(args, tokenizers)

current_epoch = Value("i", 0)

current_step = Value("i", 0)

ds_for_collator = train_dataset_group if args.max_data_loader_n_workers == 0 else None

collator = train_util.collator_class(current_epoch, current_step, ds_for_collator)

prompt_replacements = []

for emb_name in embeddings_map.keys():

emb_token_ids = embedding_to_token_ids[emb_name]

if len(emb_token_ids[0]) > 1:

token_strings = [emb_name] + [f"{emb_name}{i+1}" for i in range(len(emb_token_ids[0]) - 1)]

replace_to = " ".join(token_strings)

train_dataset_group.add_replacement(emb_name, replace_to)

prompt_replacement = (emb_name, replace_to)

prompt_replacements.append(prompt_replacement)

if args.debug_dataset:

train_util.debug_dataset(train_dataset_group)

return

if len(train_dataset_group) == 0:

logger.error(

"No data found. Please verify arguments (train_data_dir must be the parent of folders with images) / 画像がありません。引数指定を確認してください（train_data_dirには画像があるフォルダではなく、画像があるフォルダの親フォルダを指定する必要があります）"

)

return

if cache_latents:

assert (

train_dataset_group.is_latent_cacheable()

), "when caching latents, either color_aug or random_crop cannot be used / latentをキャッシュするときはcolor_augとrandom_cropは使えません"

self.assert_extra_args(args, train_dataset_group)

# 差分追加学習のためにモデルを読み込む

sys.path.append(os.path.dirname(__file__))

accelerator.print("import network module:", args.network_module)

network_module = importlib.import_module(args.network_module)

if args.base_weights is not None:

# base_weights が指定されている場合は、指定された重みを読み込みマージする

for i, weight_path in enumerate(args.base_weights):

if args.base_weights_multiplier is None or len(args.base_weights_multiplier) <= i:

multiplier = 1.0

else:

multiplier = args.base_weights_multiplier[i]

accelerator.print(f"merging module: {weight_path} with multiplier {multiplier}")

module, weights_sd = network_module.create_network_from_weights(

multiplier, weight_path, vae, text_encoder, unet, for_inference=True

)

module.merge_to(text_encoder, unet, weights_sd, weight_dtype, accelerator.device if args.lowram else "cpu")

accelerator.print(f"all weights merged: {', '.join(args.base_weights)}")

# 学習を準備する

if cache_latents:

vae.to(accelerator.device, dtype=vae_dtype)

vae.requires_grad_(False)

vae.eval()

with torch.no_grad():

train_dataset_group.cache_latents(vae, args.vae_batch_size, args.cache_latents_to_disk, accelerator.is_main_process)

vae.to("cpu")

clean_memory_on_device(accelerator.device)

accelerator.wait_for_everyone()

# move text encoder to target device, cast if not being trained

for t_enc in text_encoders:

t_enc.to(accelerator.device, dtype=(weight_dtype if not args.continue_inversion else None))

# 必要ならテキストエンコーダーの出力をキャッシュする: Text Encoderはcpuまたはgpuへ移される

# cache text encoder outputs if needed: Text Encoder is moved to cpu or gpu

# not supported and we can't set text encoder to fp16

self.cache_text_encoder_outputs_if_needed(

args, accelerator, unet, vae, tokenizers, text_encoders, train_dataset_group, weight_dtype

)

# prepare network

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 (;'∀')

if args.dim_from_weights:

network, _ = network_module.create_network_from_weights(1, args.network_weights, vae, text_encoder, unet, **net_kwargs)

else:

if "dropout" not in net_kwargs:

# workaround for LyCORIS (;^ω^)

net_kwargs["dropout"] = args.network_dropout

network = network_module.create_network(

1.0,

args.network_dim,

args.network_alpha,

vae,

text_encoder,

unet,

neuron_dropout=args.network_dropout,

**net_kwargs,

)

if network is None:

return

network_has_multiplier = hasattr(network, "set_multiplier")

if hasattr(network, "prepare_network"):

network.prepare_network(args)

if args.scale_weight_norms and not hasattr(network, "apply_max_norm_regularization"):

logger.warning(

"warning: scale_weight_norms is specified but the network does not support it / scale_weight_normsが指定されていますが、ネットワークが対応していません"

)

args.scale_weight_norms = False

train_unet = not args.network_train_text_encoder_only

train_text_encoder = self.is_train_text_encoder(args)

network.apply_to(text_encoder, unet, train_text_encoder, train_unet)

if args.network_weights is not None:

# FIXME consider alpha of weights

info = network.load_weights(args.network_weights)

accelerator.print(f"load network weights from {args.network_weights}: {info}")

# disable training if LR is 0 to save memory. useful for resuming LoRA with frozen TE or Unet

train_unet = train_unet and args.unet_lr != 0.0

train_text_encoder = train_text_encoder and args.text_encoder_lr != 0.0

if args.gradient_checkpointing:

unet.enable_gradient_checkpointing()

for t_enc in text_encoders:

t_enc.gradient_checkpointing_enable()

del t_enc

network.enable_gradient_checkpointing() # may have no effect

# 学習に必要なクラスを準備する

accelerator.print("prepare optimizer, data loader etc.")

schedulefree = "schedulefree" in args.optimizer_type.lower()

if schedulefree:

args.lr_scheduler = "constant"

if args.lr_warmup_steps:

if args.optimizer_args is None:

args.optimizer_args = []

warmup_steps = args.lr_warmup_steps

if warmup_steps < 1.0:

warmup_steps = int(warmup_steps * args.max_train_steps)

args.optimizer_args.append(f"warmup_steps={warmup_steps}")

args.lr_warmup_steps = None

# 後方互換性を確保するよ

try:

results = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr, args.learning_rate)

if type(results) is tuple:

trainable_params = results[0]

lr_descriptions = results[1]

else:

trainable_params = results

lr_descriptions = None

# Add embeddings params when continuing the inversion

if args.continue_inversion:

# TODO: might be good to add the embedding to the LoRA module directly to continue training ("bundle_emb.{emb_name}.string_to_param.*")

for t_enc in text_encoders:

trainable_params.append({

"params": t_enc.get_input_embeddings().parameters(),

"lr": args.embedding_lr or args.text_encoder_lr or args.learning_rate,

})

except TypeError as e:

# logger.warning(f"{e}")

# accelerator.print(

# "Deprecated: use prepare_optimizer_params(text_encoder_lr, unet_lr, learning_rate) instead of prepare_optimizer_params(text_encoder_lr, unet_lr)"

# )

trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr)

lr_descriptions = None

# if len(trainable_params) == 0:

# accelerator.print("no trainable parameters found / 学習可能なパラメータが見つかりませんでした")

# for params in trainable_params:

# for k, v in params.items():

# if type(v) == float:

# pass

# else:

# v = len(v)

# accelerator.print(f"trainable_params: {k} = {v}")

optimizer_name, optimizer_args, optimizer = train_util.get_optimizer(args, trainable_params)

# dataloaderを準備する

# DataLoaderのプロセス数：0 は persistent_workers が使えないので注意

n_workers = min(args.max_data_loader_n_workers, os.cpu_count()) # cpu_count or max_data_loader_n_workers

train_dataloader = torch.utils.data.DataLoader(

train_dataset_group,

batch_size=1,

shuffle=True,

collate_fn=collator,

num_workers=n_workers,

persistent_workers=args.persistent_data_loader_workers and n_workers > 0,

)

# 学習ステップ数を計算する

if args.max_train_epochs is not None:

args.max_train_steps = args.max_train_epochs * math.ceil(

len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps

)

accelerator.print(

f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}"

)

# データセット側にも学習ステップを送信

train_dataset_group.set_max_train_steps(args.max_train_steps)

# lr schedulerを用意する

lr_scheduler = train_util.get_scheduler_fix(args, optimizer, accelerator.num_processes)

# 実験的機能：勾配も含めたfp16/bf16学習を行う　モデル全体をfp16/bf16にする

if args.full_fp16:

assert (

args.mixed_precision == "fp16"

), "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"

accelerator.print("enable full fp16 training.")

network.to(weight_dtype)

elif args.full_bf16:

assert (

args.mixed_precision == "bf16"

), "full_bf16 requires mixed precision='bf16' / full_bf16を使う場合はmixed_precision='bf16'を指定してください。"

accelerator.print("enable full bf16 training.")

network.to(weight_dtype)

unet_weight_dtype = te_weight_dtype = weight_dtype

# Experimental Feature: Put base model into fp8 to save vram

if args.fp8_base:

assert torch.__version__ >= "2.1.0", "fp8_base requires torch>=2.1.0 / fp8を使う場合はtorch>=2.1.0が必要です。"

assert (

args.mixed_precision != "no"

), "fp8_base requires mixed precision='fp16' or 'bf16' / fp8を使う場合はmixed_precision='fp16'または'bf16'が必要です。"

accelerator.print("enable fp8 training.")

unet_weight_dtype = torch.float8_e4m3fn

te_weight_dtype = torch.float8_e4m3fn

unet.requires_grad_(False)

unet.to(dtype=unet_weight_dtype)

for t_enc in text_encoders:

t_enc.requires_grad_(False)

# in case of cpu, dtype is already set to fp32 because cpu does not support fp8/fp16/bf16

if t_enc.device.type != "cpu" and not args.continue_inversion:

t_enc.to(dtype=te_weight_dtype)

# nn.Embedding not support FP8

t_enc.text_model.embeddings.to(dtype=(weight_dtype if te_weight_dtype != weight_dtype else te_weight_dtype))

# acceleratorがなんかよろしくやってくれるらしい / accelerator will do something good

if args.deepspeed:

ds_model = deepspeed_utils.prepare_deepspeed_model(

args,

unet=unet if train_unet else None,

text_encoder1=text_encoders[0] if train_text_encoder else None,

text_encoder2=text_encoders[1] if train_text_encoder and len(text_encoders) > 1 else None,

network=network,

)

ds_model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(

ds_model, optimizer, train_dataloader, lr_scheduler

)

training_model = ds_model

else:

if train_unet:

unet = accelerator.prepare(unet)

else:

unet.to(accelerator.device, dtype=unet_weight_dtype) # move to device because unet is not prepared by accelerator

if train_text_encoder or args.continue_inversion:

if len(text_encoders) > 1:

text_encoder = text_encoders = [accelerator.prepare(t_enc) for t_enc in text_encoders]

else:

text_encoder = accelerator.prepare(text_encoder)

text_encoders = [text_encoder]

else:

pass # if text_encoder is not trained, no need to prepare. and device and dtype are already set

network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(

network, optimizer, train_dataloader, lr_scheduler

)

training_model = network

# Build list of original embeddings to freeze all but the modified embeddings

if args.continue_inversion:

token_ids_list = [[] for _ in text_encoders]

for sublists in embedding_to_token_ids.values():

for i, sublist in enumerate(sublists):

token_ids_list[i].extend(sublist)

index_no_updates_list = []

index_updates_list = []

orig_embeds_params_list = []

for tokenizer, token_ids, t_enc in zip(tokenizers, token_ids_list, text_encoders):

index_no_updates = torch.arange(len(tokenizer)) < token_ids[0]

index_no_updates_list.append(index_no_updates)

index_updates = ~index_no_updates

index_updates_list.append(index_updates)

orig_embeds_params = accelerator.unwrap_model(t_enc).get_input_embeddings().weight.data.detach().clone()

orig_embeds_params_list.append(orig_embeds_params)

# Freeze all parameters except for the token embeddings in text encoder

t_enc.requires_grad_(True)

unwrapped_te = accelerator.unwrap_model(t_enc)

unwrapped_te.text_model.encoder.requires_grad_(False)

unwrapped_te.text_model.final_layer_norm.requires_grad_(False)

unwrapped_te.text_model.embeddings.position_embedding.requires_grad_(False)

# unwrapped_te.text_model.embeddings.token_embedding.requires_grad_(True)

if args.gradient_checkpointing:

# according to TI example in Diffusers, train is required

unet.train()

for t_enc in text_encoders:

t_enc.train()

# set top parameter requires_grad = True for gradient checkpointing works

if train_text_encoder or args.continue_inversion:

t_enc.text_model.embeddings.requires_grad_(True)

else:

unet.eval()

for t_enc in text_encoders:

t_enc.eval()

del t_enc

accelerator.unwrap_model(network).prepare_grad_etc(text_encoder, unet)

if not train_text_encoder:

for lora in accelerator.unwrap_model(network).text_encoder_loras:

lora.requires_grad_(False)

if not cache_latents: # キャッシュしない場合はVAEを使うのでVAEを準備する

vae.requires_grad_(False)

vae.eval()

vae.to(accelerator.device, dtype=vae_dtype)

# 実験的機能：勾配も含めたfp16学習を行う　PyTorchにパッチを当ててfp16でのgrad scaleを有効にする

if args.full_fp16:

train_util.patch_accelerator_for_fp16_training(accelerator)

# before resuming make hook for saving/loading to save/load the network weights only

def save_model_hook(models, weights, output_dir):

# pop weights of other models than network to save only network weights

# only main process or deepspeed https://github.com/huggingface/diffusers/issues/2606

if accelerator.is_main_process or args.deepspeed:

remove_indices = []

for i, model in enumerate(models):

if not isinstance(model, type(accelerator.unwrap_model(network))):

remove_indices.append(i)

for i in reversed(remove_indices):

if len(weights) > i:

weights.pop(i)

# print(f"save model hook: {len(weights)} weights will be saved")

# save current ecpoch and step

train_state_file = os.path.join(output_dir, "train_state.json")

# +1 is needed because the state is saved before current_step is set from global_step

logger.info(f"save train state to {train_state_file} at epoch {current_epoch.value} step {current_step.value+1}")

with open(train_state_file, "w", encoding="utf-8") as f:

json.dump({"current_epoch": current_epoch.value, "current_step": current_step.value + 1}, f)

steps_from_state = None

def load_model_hook(models, input_dir):

# remove models except network

remove_indices = []

for i, model in enumerate(models):

if not isinstance(model, type(accelerator.unwrap_model(network))):

remove_indices.append(i)

for i in reversed(remove_indices):

models.pop(i)

# print(f"load model hook: {len(models)} models will be loaded")

# load current epoch and step to

nonlocal steps_from_state

train_state_file = os.path.join(input_dir, "train_state.json")

if os.path.exists(train_state_file):

with open(train_state_file, "r", encoding="utf-8") as f:

data = json.load(f)

steps_from_state = data["current_step"]

logger.info(f"load train state from {train_state_file}: {data}")

accelerator.register_save_state_pre_hook(save_model_hook)

accelerator.register_load_state_pre_hook(load_model_hook)

# resumeする

train_util.resume_from_local_or_hf_if_specified(accelerator, args)

# 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

# 学習する

# TODO: find a way to handle total batch size when there are multiple datasets

total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps

accelerator.print("running training / 学習開始")

accelerator.print(f" num train images * repeats / 学習画像の数×繰り返し回数: {train_dataset_group.num_train_images}")

accelerator.print(f" num reg images / 正則化画像の数: {train_dataset_group.num_reg_images}")

accelerator.print(f" num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}")

accelerator.print(f" num epochs / epoch数: {num_train_epochs}")

accelerator.print(

f" batch size per device / バッチサイズ: {', '.join([str(d.batch_size) for d in train_dataset_group.datasets])}"

)

# accelerator.print(f" total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}")

accelerator.print(f" gradient accumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}")

accelerator.print(f" total optimization steps / 学習ステップ数: {args.max_train_steps}")

# TODO refactor metadata creation and move to util

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_group.num_train_images,

"ss_num_reg_images": train_dataset_group.num_reg_images,

"ss_num_batches_per_epoch": len(train_dataloader),

"ss_num_epochs": num_train_epochs,

"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": 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 have alpha

"ss_network_dropout": args.network_dropout, # some networks may not have dropout

"ss_mixed_precision": args.mixed_precision,

"ss_full_fp16": bool(args.full_fp16),

"ss_v2": bool(args.v2),

"ss_base_model_version": model_version,

"ss_clip_skip": args.clip_skip,

"ss_max_token_length": args.max_token_length,

"ss_cache_latents": bool(args.cache_latents),

"ss_seed": args.seed,

"ss_lowram": args.lowram,

"ss_noise_offset": args.noise_offset,

"ss_multires_noise_iterations": args.multires_noise_iterations,

"ss_multires_noise_discount": args.multires_noise_discount,

"ss_adaptive_noise_scale": args.adaptive_noise_scale,

"ss_zero_terminal_snr": args.zero_terminal_snr,

"ss_training_comment": args.training_comment, # will not be updated after training

"ss_sd_scripts_commit_hash": train_util.get_git_revision_hash(),

"ss_optimizer": optimizer_name + (f"({optimizer_args})" if len(optimizer_args) > 0 else ""),

"ss_max_grad_norm": args.max_grad_norm,

"ss_caption_dropout_rate": args.caption_dropout_rate,

"ss_caption_dropout_every_n_epochs": args.caption_dropout_every_n_epochs,

"ss_caption_tag_dropout_rate": args.caption_tag_dropout_rate,

"ss_face_crop_aug_range": args.face_crop_aug_range,

"ss_prior_loss_weight": args.prior_loss_weight,

"ss_min_snr_gamma": args.min_snr_gamma,

"ss_scale_weight_norms": args.scale_weight_norms,

"ss_ip_noise_gamma": args.ip_noise_gamma,

"ss_debiased_estimation": bool(args.debiased_estimation_loss),

"ss_noise_offset_random_strength": args.noise_offset_random_strength,

"ss_ip_noise_gamma_random_strength": args.ip_noise_gamma_random_strength,

"ss_loss_type": args.loss_type,

"ss_huber_schedule": args.huber_schedule,

"ss_huber_c": args.huber_c,

}

if use_user_config:

# save metadata of multiple datasets

# NOTE: pack "ss_datasets" value as json one time

# or should also pack nested collections as json?

datasets_metadata = []

tag_frequency = {} # merge tag frequency for metadata editor

dataset_dirs_info = {} # merge subset dirs for metadata editor

for dataset in train_dataset_group.datasets:

is_dreambooth_dataset = isinstance(dataset, DreamBoothDataset)

dataset_metadata = {

"is_dreambooth": is_dreambooth_dataset,

"batch_size_per_device": dataset.batch_size,

"num_train_images": dataset.num_train_images, # includes repeating

"num_reg_images": dataset.num_reg_images,

"resolution": (dataset.width, dataset.height),

"enable_bucket": bool(dataset.enable_bucket),

"min_bucket_reso": dataset.min_bucket_reso,

"max_bucket_reso": dataset.max_bucket_reso,

"tag_frequency": dataset.tag_frequency,

"bucket_info": dataset.bucket_info,

}

subsets_metadata = []

for subset in dataset.subsets:

subset_metadata = {

"img_count": subset.img_count,

"num_repeats": subset.num_repeats,

"color_aug": bool(subset.color_aug),

"flip_aug": bool(subset.flip_aug),

"random_crop": bool(subset.random_crop),

"shuffle_caption": bool(subset.shuffle_caption),

"keep_tokens": subset.keep_tokens,

"keep_tokens_separator": subset.keep_tokens_separator,

"secondary_separator": subset.secondary_separator,

"enable_wildcard": bool(subset.enable_wildcard),

"caption_prefix": subset.caption_prefix,

"caption_suffix": subset.caption_suffix,

}

image_dir_or_metadata_file = None

if subset.image_dir:

image_dir = os.path.basename(subset.image_dir)

subset_metadata["image_dir"] = image_dir

image_dir_or_metadata_file = image_dir

if is_dreambooth_dataset:

subset_metadata["class_tokens"] = subset.class_tokens

subset_metadata["is_reg"] = subset.is_reg

if subset.is_reg:

image_dir_or_metadata_file = None # not merging reg dataset

else:

metadata_file = os.path.basename(subset.metadata_file)

subset_metadata["metadata_file"] = metadata_file

image_dir_or_metadata_file = metadata_file # may overwrite

subsets_metadata.append(subset_metadata)

# merge dataset dir: not reg subset only

# TODO update additional-network extension to show detailed dataset config from metadata

if image_dir_or_metadata_file is not None:

# datasets may have a certain dir multiple times

v = image_dir_or_metadata_file

i = 2

while v in dataset_dirs_info:

v = image_dir_or_metadata_file + f" ({i})"

i += 1

image_dir_or_metadata_file = v

dataset_dirs_info[image_dir_or_metadata_file] = {

"n_repeats": subset.num_repeats,

"img_count": subset.img_count,

}

dataset_metadata["subsets"] = subsets_metadata

datasets_metadata.append(dataset_metadata)

# merge tag frequency:

for ds_dir_name, ds_freq_for_dir in dataset.tag_frequency.items():

# あるディレクトリが複数のdatasetで使用されている場合、一度だけ数える

# もともと繰り返し回数を指定しているので、キャプション内でのタグの出現回数と、それが学習で何度使われるかは一致しない

# なので、ここで複数datasetの回数を合算してもあまり意味はない

if ds_dir_name in tag_frequency:

continue

tag_frequency[ds_dir_name] = ds_freq_for_dir

metadata["ss_datasets"] = json.dumps(datasets_metadata)

metadata["ss_tag_frequency"] = json.dumps(tag_frequency)

metadata["ss_dataset_dirs"] = json.dumps(dataset_dirs_info)

else:

# conserving backward compatibility when using train_dataset_dir and reg_dataset_dir

assert (

len(train_dataset_group.datasets) == 1

), f"There should be a single dataset but {len(train_dataset_group.datasets)} found. This seems to be a bug. / データセットは1個だけ存在するはずですが、実際には{len(train_dataset_group.datasets)}個でした。プログラムのバグかもしれません。"

dataset = train_dataset_group.datasets[0]

dataset_dirs_info = {}

reg_dataset_dirs_info = {}

if use_dreambooth_method:

for subset in dataset.subsets:

info = reg_dataset_dirs_info if subset.is_reg else dataset_dirs_info

info[os.path.basename(subset.image_dir)] = {"n_repeats": subset.num_repeats, "img_count": subset.img_count}

else:

for subset in dataset.subsets:

dataset_dirs_info[os.path.basename(subset.metadata_file)] = {

"n_repeats": subset.num_repeats,

"img_count": subset.img_count,

}

metadata.update(

{

"ss_batch_size_per_device": args.train_batch_size,

"ss_total_batch_size": total_batch_size,

"ss_resolution": args.resolution,

"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_enable_bucket": bool(dataset.enable_bucket),

"ss_bucket_no_upscale": bool(dataset.bucket_no_upscale),

"ss_min_bucket_reso": dataset.min_bucket_reso,

"ss_max_bucket_reso": dataset.max_bucket_reso,

"ss_keep_tokens": args.keep_tokens,

"ss_dataset_dirs": json.dumps(dataset_dirs_info),

"ss_reg_dataset_dirs": json.dumps(reg_dataset_dirs_info),

"ss_tag_frequency": json.dumps(dataset.tag_frequency),

"ss_bucket_info": json.dumps(dataset.bucket_info),

}

)

# add extra args

if args.network_args:

metadata["ss_network_args"] = json.dumps(net_kwargs)

# model name and hash

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

if args.todo_factor:

metadata["ss_todo_factor"] = args.todo_factor

metadata["ss_todo_max_depth"] = args.todo_max_depth

metadata = {k: str(v) for k, v in metadata.items()}

# make minimum metadata for filtering

minimum_metadata = {}

for key in train_util.SS_METADATA_MINIMUM_KEYS:

if key in metadata:

minimum_metadata[key] = metadata[key]

# calculate steps to skip when resuming or starting from a specific step

initial_step = 0

if args.initial_epoch is not None or args.initial_step is not None:

# if initial_epoch or initial_step is specified, steps_from_state is ignored even when resuming

if steps_from_state is not None:

logger.warning(

"steps from the state is ignored because initial_step is specified / initial_stepが指定されているため、stateからのステップ数は無視されます"

)

if args.initial_step is not None:

initial_step = args.initial_step

else:

# num steps per epoch is calculated by num_processes and gradient_accumulation_steps

initial_step = (args.initial_epoch - 1) * math.ceil(

len(train_dataloader) / accelerator.num_processes / args.gradient_accumulation_steps

)

else:

# if initial_epoch and initial_step are not specified, steps_from_state is used when resuming

if steps_from_state is not None:

initial_step = steps_from_state

steps_from_state = None

if initial_step > 0:

assert (

args.max_train_steps > initial_step

), f"max_train_steps should be greater than initial step / max_train_stepsは初期ステップより大きい必要があります: {args.max_train_steps} vs {initial_step}"

progress_bar = tqdm(

initial=initial_step, total=args.max_train_steps, smoothing=0.1, disable=not accelerator.is_local_main_process, desc="steps"

)

epoch_to_start = 0

if initial_step > 0:

if args.skip_until_initial_step:

# if skip_until_initial_step is specified, load data and discard it to ensure the same data is used

if not args.resume:

logger.info(

f"initial_step is specified but not resuming. lr scheduler will be started from the beginning / initial_stepが指定されていますがresumeしていないため、lr schedulerは最初から始まります"

)

logger.info(f"skipping {initial_step} steps / {initial_step}ステップをスキップします")

# initial_step *= args.gradient_accumulation_steps

# set epoch to start to make initial_step less than len(train_dataloader)

epoch_to_start = initial_step // math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)