1 files changed, 50 insertions, 11 deletions
diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index b915b091..5a7be422 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -11,9 +11,11 @@ import datetime
 import csv
 import safetensors.torch
 
+import numpy as np
 from PIL import Image, PngImagePlugin
+from torch.utils.tensorboard import SummaryWriter
 
-from modules import shared, devices, sd_hijack, processing, sd_models, images, sd_samplers
+from modules import shared, devices, sd_hijack, processing, sd_models, images, sd_samplers, sd_hijack_checkpoint
 import modules.textual_inversion.dataset
 from modules.textual_inversion.learn_schedule import LearnRateScheduler
 
@@ -248,11 +250,14 @@ def create_embedding(name, num_vectors_per_token, overwrite_old, init_text='*'):
     with devices.autocast():
         cond_model([""])  # will send cond model to GPU if lowvram/medvram is active
 
-    embedded = cond_model.encode_embedding_init_text(init_text, num_vectors_per_token)
+    #cond_model expects at least some text, so we provide '*' as backup.
+    embedded = cond_model.encode_embedding_init_text(init_text or '*', num_vectors_per_token)
     vec = torch.zeros((num_vectors_per_token, embedded.shape[1]), device=devices.device)
 
-    for i in range(num_vectors_per_token):
-        vec[i] = embedded[i * int(embedded.shape[0]) // num_vectors_per_token]
+    #Only copy if we provided an init_text, otherwise keep vectors as zeros
+    if init_text:
+        for i in range(num_vectors_per_token):
+            vec[i] = embedded[i * int(embedded.shape[0]) // num_vectors_per_token]
 
     # Remove illegal characters from name.
     name = "".join( x for x in name if (x.isalnum() or x in "._- "))
@@ -291,6 +296,30 @@ def write_loss(log_directory, filename, step, epoch_len, values):
             **values,
         })
 
+def tensorboard_setup(log_directory):
+    os.makedirs(os.path.join(log_directory, "tensorboard"), exist_ok=True)
+    return SummaryWriter(
+            log_dir=os.path.join(log_directory, "tensorboard"),
+            flush_secs=shared.opts.training_tensorboard_flush_every)
+
+def tensorboard_add(tensorboard_writer, loss, global_step, step, learn_rate, epoch_num):
+    tensorboard_add_scaler(tensorboard_writer, "Loss/train", loss, global_step)
+    tensorboard_add_scaler(tensorboard_writer, f"Loss/train/epoch-{epoch_num}", loss, step)
+    tensorboard_add_scaler(tensorboard_writer, "Learn rate/train", learn_rate, global_step)
+    tensorboard_add_scaler(tensorboard_writer, f"Learn rate/train/epoch-{epoch_num}", learn_rate, step)
+
+def tensorboard_add_scaler(tensorboard_writer, tag, value, step):
+    tensorboard_writer.add_scalar(tag=tag, 
+        scalar_value=value, global_step=step)
+
+def tensorboard_add_image(tensorboard_writer, tag, pil_image, step):
+    # Convert a pil image to a torch tensor
+    img_tensor = torch.as_tensor(np.array(pil_image, copy=True))
+    img_tensor = img_tensor.view(pil_image.size[1], pil_image.size[0], 
+        len(pil_image.getbands()))
+    img_tensor = img_tensor.permute((2, 0, 1))
+                
+    tensorboard_writer.add_image(tag, img_tensor, global_step=step)
 
 def validate_train_inputs(model_name, learn_rate, batch_size, gradient_step, data_root, template_file, template_filename, steps, save_model_every, create_image_every, log_directory, name="embedding"):
     assert model_name, f"{name} not selected"
@@ -316,7 +345,7 @@ def validate_train_inputs(model_name, learn_rate, batch_size, gradient_step, dat
         assert log_directory, "Log directory is empty"
 
 
-def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_root, log_directory, training_width, training_height, varsize, steps, clip_grad_mode, clip_grad_value, shuffle_tags, tag_drop_out, latent_sampling_method, create_image_every, save_embedding_every, template_filename, 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):
+def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_step, data_root, log_directory, training_width, training_height, varsize, steps, clip_grad_mode, clip_grad_value, shuffle_tags, tag_drop_out, latent_sampling_method, create_image_every, save_embedding_every, template_filename, 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):
     save_embedding_every = save_embedding_every or 0
     create_image_every = create_image_every or 0
     template_file = textual_inversion_templates.get(template_filename, None)
@@ -369,13 +398,16 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
     # 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
+    
+    if shared.opts.training_enable_tensorboard:
+        tensorboard_writer = tensorboard_setup(log_directory)
 
     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=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, varsize=varsize)
 
     if shared.opts.save_training_settings_to_txt:
-        save_settings_to_file(log_directory, {**dict(model_name=checkpoint.model_name, model_hash=checkpoint.hash, num_of_dataset_images=len(ds), num_vectors_per_token=len(embedding.vec)), **locals()})
+        save_settings_to_file(log_directory, {**dict(model_name=checkpoint.model_name, model_hash=checkpoint.shorthash, num_of_dataset_images=len(ds), num_vectors_per_token=len(embedding.vec)), **locals()})
 
     latent_sampling_method = ds.latent_sampling_method
 
@@ -420,6 +452,8 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
 
     pbar = tqdm.tqdm(total=steps - initial_step)
     try:
+        sd_hijack_checkpoint.add()
+
         for i in range((steps-initial_step) * gradient_step):
             if scheduler.finished:
                 break
@@ -476,9 +510,8 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
                 epoch_num = embedding.step // steps_per_epoch
                 epoch_step = embedding.step % steps_per_epoch
 
-                description = f"Training textual inversion [Epoch {epoch_num}: {epoch_step+1}/{steps_per_epoch}]loss: {loss_step:.7f}"
+                description = f"Training textual inversion [Epoch {epoch_num}: {epoch_step+1}/{steps_per_epoch}] loss: {loss_step:.7f}"
                 pbar.set_description(description)
-                shared.state.textinfo = description
                 if embedding_dir is not None and steps_done % save_embedding_every == 0:
                     # Before saving, change name to match current checkpoint.
                     embedding_name_every = f'{embedding_name}-{steps_done}'
@@ -528,10 +561,14 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
                         shared.sd_model.first_stage_model.to(devices.cpu)
 
                     if image is not None:
-                        shared.state.current_image = image
+                        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')
@@ -549,7 +586,7 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
 
                         checkpoint = sd_models.select_checkpoint()
                         footer_left = checkpoint.model_name
-                        footer_mid = '[{}]'.format(checkpoint.hash)
+                        footer_mid = '[{}]'.format(checkpoint.shorthash)
                         footer_right = '{}v {}s'.format(vectorSize, steps_done)
 
                         captioned_image = caption_image_overlay(image, title, footer_left, footer_mid, footer_right)
@@ -582,16 +619,18 @@ Last saved image: {html.escape(last_saved_image)}<br/>
         pbar.close()
         shared.sd_model.first_stage_model.to(devices.device)
         shared.parallel_processing_allowed = old_parallel_processing_allowed
+        sd_hijack_checkpoint.remove()
 
     return embedding, filename
 
+
 def save_embedding(embedding, optimizer, checkpoint, embedding_name, filename, remove_cached_checksum=True):
     old_embedding_name = embedding.name
     old_sd_checkpoint = embedding.sd_checkpoint if hasattr(embedding, "sd_checkpoint") else None
     old_sd_checkpoint_name = embedding.sd_checkpoint_name if hasattr(embedding, "sd_checkpoint_name") else None
     old_cached_checksum = embedding.cached_checksum if hasattr(embedding, "cached_checksum") else None
     try:
-        embedding.sd_checkpoint = checkpoint.hash
+        embedding.sd_checkpoint = checkpoint.shorthash
         embedding.sd_checkpoint_name = checkpoint.model_name
         if remove_cached_checksum:
             embedding.cached_checksum = None