import csv import datetime import gc import html import os import sys import traceback import torch import tqdm from PIL import PngImagePlugin from .dataset import PersonalizedBase, PersonalizedDataLoader from ..scheduler import CosineAnnealingWarmUpRestarts from ..hnutil import optim_to from modules import shared, devices, sd_models, images, processing, sd_samplers, sd_hijack from modules.textual_inversion.image_embedding import caption_image_overlay, insert_image_data_embed, embedding_to_b64 from modules.textual_inversion.learn_schedule import LearnRateScheduler from modules.textual_inversion.textual_inversion import save_embedding from torch.utils.tensorboard import SummaryWriter from modules.textual_inversion.textual_inversion import tensorboard_add, tensorboard_setup, tensorboard_add_scaler, tensorboard_add_image #apply OsError avoid here delayed_values = {} def write_loss(log_directory, filename, step, epoch_len, values): if shared.opts.training_write_csv_every == 0: return if step % shared.opts.training_write_csv_every != 0: return write_csv_header = False if os.path.exists(os.path.join(log_directory, filename)) else True try: with open(os.path.join(log_directory, filename), "a+", newline='') as fout: csv_writer = csv.DictWriter(fout, fieldnames=["step", "epoch", "epoch_step", *(values.keys())]) if write_csv_header: csv_writer.writeheader() if log_directory + filename in delayed_values: delayed = delayed_values[log_directory + filename] for step, epoch, epoch_step, values in delayed: csv_writer.writerow({ "step": step, "epoch": epoch, "epoch_step": epoch_step, **values, }) delayed.clear() epoch, epoch_step = divmod(step - 1, epoch_len) csv_writer.writerow({ "step": step, "epoch": epoch, "epoch_step": epoch_step, **values, }) except OSError: epoch, epoch_step = divmod(step-1, epoch_len) if log_directory + filename in delayed_values: delayed_values[log_directory + filename].append((step , epoch, epoch_step, values)) else: delayed_values[log_directory + filename] = [(step, epoch, epoch_step, values)] def validate_train_inputs(model_name, learn_rate, batch_size, gradient_step, data_root, template_file, steps, save_model_every, create_image_every, log_directory, name="embedding"): assert model_name, f"{name} not selected" assert learn_rate, "Learning rate is empty or 0" assert isinstance(batch_size, int), "Batch size must be integer" assert batch_size > 0, "Batch size must be positive" assert isinstance(gradient_step, int), "Gradient accumulation step must be integer" assert gradient_step > 0, "Gradient accumulation step must be positive" assert data_root, "Dataset directory is empty" assert os.path.isdir(data_root), "Dataset directory doesn't exist" assert os.listdir(data_root), "Dataset directory is empty" assert template_file, "Prompt template file is empty" assert os.path.isfile(template_file), "Prompt template file doesn't exist" assert steps, "Max steps is empty or 0" assert isinstance(steps, int), "Max steps must be integer" assert steps > 0, "Max steps must be positive" assert isinstance(save_model_every, int), "Save {name} must be integer" assert save_model_every >= 0, "Save {name} must be positive or 0" assert isinstance(create_image_every, int), "Create image must be integer" assert create_image_every >= 0, "Create image must be positive or 0" if save_model_every or create_image_every: assert log_directory, "Log directory is empty" def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_step, data_root, log_directory, training_width, training_height, steps, shuffle_tags, tag_drop_out, latent_sampling_method, create_image_every, save_embedding_every, template_file, save_image_with_stored_embedding, preview_from_txt2img, preview_prompt, preview_negative_prompt, preview_steps, preview_sampler_index, preview_cfg_scale, preview_seed, preview_width, preview_height, use_beta_scheduler=False, beta_repeat_epoch=4000, epoch_mult=1,warmup =10, min_lr=1e-7, gamma_rate=1, save_when_converge=False, create_when_converge=False, move_optimizer=True, use_adamw_parameter=False, adamw_weight_decay=0.01, adamw_beta_1=0.9, adamw_beta_2=0.99,adamw_eps=1e-8, use_grad_opts=False, gradient_clip_opt='None', optional_gradient_clip_value=1e01, optional_gradient_norm_type=2 ): save_embedding_every = save_embedding_every or 0 create_image_every = create_image_every or 0 validate_train_inputs(embedding_name, learn_rate, batch_size, gradient_step, data_root, template_file, steps, save_embedding_every, create_image_every, log_directory, name="embedding") try: if use_adamw_parameter: adamw_weight_decay, adamw_beta_1, adamw_beta_2, adamw_eps = [float(x) for x in [adamw_weight_decay, adamw_beta_1, adamw_beta_2, adamw_eps]] assert 0 <= adamw_weight_decay, "Weight decay paramter should be larger or equal than zero!" assert (all(0 <= x <= 1 for x in [adamw_beta_1, adamw_beta_2, adamw_eps])), "Cannot use negative or >1 number for adamW parameters!" adamW_kwarg_dict = { 'weight_decay' : adamw_weight_decay, 'betas' : (adamw_beta_1, adamw_beta_2), 'eps' : adamw_eps } print('Using custom AdamW parameters') else: adamW_kwarg_dict = { 'weight_decay' : 0.01, 'betas' : (0.9, 0.99), 'eps' : 1e-8 } if use_beta_scheduler: print("Using Beta Scheduler") beta_repeat_epoch = int(beta_repeat_epoch) assert beta_repeat_epoch > 0, f"Cannot use too small cycle {beta_repeat_epoch}!" min_lr = float(min_lr) assert min_lr < 1, f"Cannot use minimum lr with {min_lr}!" gamma_rate = float(gamma_rate) print(f"Using learn rate decay(per cycle) of {gamma_rate}") assert 0 <= gamma_rate <= 1, f"Cannot use gamma rate with {gamma_rate}!" epoch_mult = float(epoch_mult) assert 1 <= epoch_mult, "Cannot use epoch multiplier smaller than 1!" warmup = int(warmup) assert warmup >= 1, "Warmup epoch should be larger than 0!" print(f"Save when converges : {save_when_converge}") print(f"Generate image when converges : {create_when_converge}") else: beta_repeat_epoch = 4000 epoch_mult=1 warmup=10 min_lr=1e-7 gamma_rate=1 save_when_converge = False create_when_converge = False except ValueError: raise RuntimeError("Cannot use advanced LR scheduler settings!") if use_grad_opts and gradient_clip_opt != "None": try: optional_gradient_clip_value = float(optional_gradient_clip_value) except ValueError: raise RuntimeError(f"Cannot convert invalid gradient clipping value {optional_gradient_clip_value})") if gradient_clip_opt == "Norm": try: grad_norm = int(optional_gradient_norm_type) except ValueError: raise RuntimeError(f"Cannot convert invalid gradient norm type {optional_gradient_norm_type})") assert grad_norm >= 0, f"P-norm cannot be calculated from negative number {grad_norm}" def gradient_clipping(arg1): torch.nn.utils.clip_grad_norm_(arg1, optional_gradient_clip_value, optional_gradient_norm_type) return else: def gradient_clipping(arg1): torch.nn.utils.clip_grad_value_(arg1, optional_gradient_clip_value) return else: def gradient_clipping(arg1): return # Function gradient clipping is inplace(_) operation. shared.state.textinfo = "Initializing textual inversion training..." shared.state.job_count = steps filename = os.path.join(shared.cmd_opts.embeddings_dir, f'{embedding_name}.pt') log_directory = os.path.join(log_directory, datetime.datetime.now().strftime("%Y-%m-%d"), embedding_name) unload = shared.opts.unload_models_when_training if save_embedding_every > 0 or save_when_converge: embedding_dir = os.path.join(log_directory, "embeddings") os.makedirs(embedding_dir, exist_ok=True) else: embedding_dir = None if create_image_every > 0 or create_when_converge: images_dir = os.path.join(log_directory, "images") os.makedirs(images_dir, exist_ok=True) else: images_dir = None if (create_image_every > 0 or create_when_converge) and save_image_with_stored_embedding: images_embeds_dir = os.path.join(log_directory, "image_embeddings") os.makedirs(images_embeds_dir, exist_ok=True) else: images_embeds_dir = None hijack = sd_hijack.model_hijack embedding = hijack.embedding_db.word_embeddings[embedding_name] checkpoint = sd_models.select_checkpoint() initial_step = embedding.step or 0 if initial_step >= steps: shared.state.textinfo = f"Model has already been trained beyond specified max steps" return embedding, filename scheduler = LearnRateScheduler(learn_rate, steps, initial_step) # 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)}..." old_parallel_processing_allowed = shared.parallel_processing_allowed tensorboard_writer = None if shared.opts.training_enable_tensorboard: print("Tensorboard logging enabled") tensorboard_writer = tensorboard_setup(log_directory) pin_memory = shared.opts.pin_memory shared.sd_model.cond_stage_model.to(devices.device) shared.sd_model.first_stage_model.to(devices.device) ds = PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=embedding_name, model=shared.sd_model, cond_model=shared.sd_model.cond_stage_model, device=devices.device, template_file=template_file, batch_size=batch_size, gradient_step=gradient_step, shuffle_tags=shuffle_tags, tag_drop_out=tag_drop_out, latent_sampling_method=latent_sampling_method) latent_sampling_method = ds.latent_sampling_method dl = PersonalizedDataLoader(ds, latent_sampling_method=latent_sampling_method, batch_size=ds.batch_size, pin_memory=pin_memory) if unload: shared.parallel_processing_allowed = False shared.sd_model.first_stage_model.to(devices.cpu) embedding.vec.requires_grad = True optimizer_name = 'AdamW' # hardcoded optimizer name now if use_adamw_parameter: optimizer = torch.optim.AdamW(params=[embedding.vec], lr=scheduler.learn_rate, **adamW_kwarg_dict) else: optimizer = torch.optim.AdamW([embedding.vec], lr=scheduler.learn_rate, weight_decay=0.0) if os.path.exists(filename + '.optim'): # This line must be changed if Optimizer type can be different from saved optimizer. try: optimizer_saved_dict = torch.load(filename + '.optim', map_location='cpu') if embedding.checksum() == optimizer_saved_dict.get('hash', None): optimizer_state_dict = optimizer_saved_dict.get('optimizer_state_dict', None) optimizer.load_state_dict(optimizer_state_dict) print("Loaded existing optimizer from checkpoint") except RuntimeError as e: print("Cannot resume from saved optimizer!") print(e) else: print("No saved optimizer exists in checkpoint") optim_to(optimizer, devices.device) if use_beta_scheduler: scheduler_beta = CosineAnnealingWarmUpRestarts(optimizer=optimizer, first_cycle_steps=beta_repeat_epoch, cycle_mult=epoch_mult, max_lr=scheduler.learn_rate, warmup_steps=warmup, min_lr=min_lr, gamma=gamma_rate) scheduler_beta.last_epoch = embedding.step-1 else: scheduler_beta = None for pg in optimizer.param_groups: pg['lr'] = scheduler.learn_rate 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 last_saved_file = "" last_saved_image = "" forced_filename = "" embedding_yet_to_be_embedded = False is_training_inpainting_model = shared.sd_model.model.conditioning_key in {'hybrid', 'concat'} img_c = None pbar = tqdm.tqdm(total=steps - initial_step) try: for i 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 if use_beta_scheduler: scheduler_beta.step(embedding.step) else: scheduler.apply(optimizer, embedding.step) if scheduler.finished: break if shared.state.interrupted: break with torch.autocast("cuda"): # c = stack_conds(batch.cond).to(devices.device) # mask = torch.tensor(batch.emb_index).to(devices.device, non_blocking=pin_memory) # print(mask) # c[:, 1:1+embedding.vec.shape[0]] = embedding.vec.to(devices.device, non_blocking=pin_memory) x = batch.latent_sample.to(devices.device, non_blocking=pin_memory) c = shared.sd_model.cond_stage_model(batch.cond_text) if is_training_inpainting_model: if img_c is None: img_c = processing.txt2img_image_conditioning(shared.sd_model, c, training_width, training_height) cond = {"c_concat": [img_c], "c_crossattn": [c]} else: cond = c loss = shared.sd_model(x, cond)[0] / gradient_step del x _loss_step += loss.item() scaler.scale(loss).backward() # go back until we reach gradient accumulation steps if (j + 1) % gradient_step != 0: continue gradient_clipping(embedding.vec) scaler.step(optimizer) scaler.update() embedding.step += 1 pbar.update() optimizer.zero_grad(set_to_none=True) loss_step = _loss_step _loss_step = 0 steps_done = embedding.step + 1 epoch_num = embedding.step // steps_per_epoch epoch_step = embedding.step % steps_per_epoch pbar.set_description(f"[Epoch {epoch_num}: {epoch_step + 1}/{steps_per_epoch}]loss: {loss_step:.7f}") if embedding_dir is not None and ((use_beta_scheduler and scheduler_beta.is_EOC(embedding.step) and save_when_converge) or (save_embedding_every > 0 and steps_done % save_embedding_every == 0)): # Before saving, change name to match current checkpoint. embedding_name_every = f'{embedding_name}-{steps_done}' last_saved_file = os.path.join(embedding_dir, f'{embedding_name_every}.pt') # if shared.opts.save_optimizer_state: # embedding.optimizer_state_dict = optimizer.state_dict() save_embedding(embedding, optimizer, checkpoint, embedding_name_every, last_saved_file, remove_cached_checksum=True) embedding_yet_to_be_embedded = True write_loss(log_directory, "textual_inversion_loss.csv", embedding.step, steps_per_epoch, { "loss": f"{loss_step:.7f}", "learn_rate": scheduler.learn_rate }) if images_dir is not None and ((use_beta_scheduler and scheduler_beta.is_EOC(embedding.step) and create_when_converge) or (create_image_every > 0 and steps_done % create_image_every == 0)): forced_filename = f'{embedding_name}-{steps_done}' last_saved_image = os.path.join(images_dir, forced_filename) rng_state = torch.get_rng_state() cuda_rng_state = None if torch.cuda.is_available(): cuda_rng_state = torch.cuda.get_rng_state_all() if move_optimizer: optim_to(optimizer, devices.cpu) gc.collect() 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, do_not_reload_embeddings=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[preview_sampler_index].name 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 processed = processing.process_images(p) image = processed.images[0] if len(processed.images) > 0 else None if move_optimizer: optim_to(optimizer, devices.device) if image is not None: shared.state.assign_current_image(image) 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}" if shared.opts.training_enable_tensorboard and shared.opts.training_tensorboard_save_images: tensorboard_add_image(tensorboard_writer, f"Validation at epoch {epoch_num}", image, embedding.step) if save_image_with_stored_embedding and os.path.exists( last_saved_file) and embedding_yet_to_be_embedded: last_saved_image_chunks = os.path.join(images_embeds_dir, f'{embedding_name}-{steps_done}.png') info = PngImagePlugin.PngInfo() data = torch.load(last_saved_file) info.add_text("sd-ti-embedding", embedding_to_b64(data)) title = "<{}>".format(data.get('name', '???')) try: vectorSize = list(data['string_to_param'].values())[0].shape[0] except Exception as e: vectorSize = '?' checkpoint = sd_models.select_checkpoint() footer_left = checkpoint.model_name footer_mid = '[{}]'.format(checkpoint.shorthash if hasattr(checkpoint, 'shorthash') else checkpoint.hash) footer_right = '{}v {}s'.format(vectorSize, steps_done) captioned_image = caption_image_overlay(image, title, footer_left, footer_mid, footer_right) captioned_image = insert_image_data_embed(captioned_image, data) captioned_image.save(last_saved_image_chunks, "PNG", pnginfo=info) embedding_yet_to_be_embedded = False if unload: 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) 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 = embedding.step shared.state.textinfo = f"""

Loss: {loss_step:.7f}
Step: {steps_done}
Last prompt: {html.escape(batch.cond_text[0])}
Last saved embedding: {html.escape(last_saved_file)}
Last saved image: {html.escape(last_saved_image)}

""" filename = os.path.join(shared.cmd_opts.embeddings_dir, f'{embedding_name}.pt') save_embedding(embedding, optimizer, checkpoint, embedding_name, filename, remove_cached_checksum=True) except Exception: print(traceback.format_exc(), file=sys.stderr) pass finally: pbar.leave = False pbar.close() shared.sd_model.first_stage_model.to(devices.device) shared.parallel_processing_allowed = old_parallel_processing_allowed return embedding, filename