1 files changed, 58 insertions, 14 deletions

diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index fd253477..71e07bcc 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -28,6 +28,7 @@ class Embedding:
         self.cached_checksum = None

         self.sd_checkpoint = None

         self.sd_checkpoint_name = None

+        self.optimizer_state_dict = None

 

     def save(self, filename):

         embedding_data = {

@@ -41,6 +42,13 @@ class Embedding:
 

         torch.save(embedding_data, filename)

 

+        if shared.opts.save_optimizer_state and self.optimizer_state_dict is not None:

+            optimizer_saved_dict = {

+                'hash': self.checksum(),

+                'optimizer_state_dict': self.optimizer_state_dict,

+            }

+            torch.save(optimizer_saved_dict, filename + '.optim')

+

     def checksum(self):

         if self.cached_checksum is not None:

             return self.cached_checksum

@@ -95,9 +103,10 @@ class EmbeddingDatabase:
         self.expected_shape = self.get_expected_shape()

 

         def process_file(path, filename):

-            name = os.path.splitext(filename)[0]

+            name, ext = os.path.splitext(filename)

+            ext = ext.upper()

 

-            if os.path.splitext(filename.upper())[-1] in ['.PNG', '.WEBP', '.JXL', '.AVIF']:

+            if ext in ['.PNG', '.WEBP', '.JXL', '.AVIF']:

                 embed_image = Image.open(path)

                 if hasattr(embed_image, 'text') and 'sd-ti-embedding' in embed_image.text:

                     data = embedding_from_b64(embed_image.text['sd-ti-embedding'])

@@ -105,8 +114,10 @@ class EmbeddingDatabase:
                 else:

                     data = extract_image_data_embed(embed_image)

                     name = data.get('name', name)

-            else:

+            elif ext in ['.BIN', '.PT']:

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

+            else:

+                return

 

             # textual inversion embeddings

             if 'string_to_param' in data:

@@ -240,11 +251,12 @@ def validate_train_inputs(model_name, learn_rate, batch_size, gradient_step, dat
     if save_model_every or create_image_every:

         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, 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):

+def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_root, log_directory, training_width, training_height, steps, clip_grad_mode, clip_grad_value, 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):

     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")

 

+    shared.state.job = "train-embedding"

     shared.state.textinfo = "Initializing textual inversion training..."

     shared.state.job_count = steps

 

@@ -282,6 +294,11 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
         return embedding, 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, ititial_step, verbose=False)

     # 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

@@ -300,6 +317,19 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
 

     embedding.vec.requires_grad = True

     optimizer = torch.optim.AdamW([embedding.vec], lr=scheduler.learn_rate, weight_decay=0.0)

+    if shared.opts.save_optimizer_state:

+        optimizer_state_dict = None

+        if os.path.exists(filename + '.optim'):

+            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)

+    

+        if optimizer_state_dict is not None:

+            optimizer.load_state_dict(optimizer_state_dict)

+            print("Loaded existing optimizer from checkpoint")

+        else:

+            print("No saved optimizer exists in checkpoint")

+

     scaler = torch.cuda.amp.GradScaler()

 

     batch_size = ds.batch_size

@@ -315,6 +345,9 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
     forced_filename = "<none>"

     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):

@@ -332,14 +365,22 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
                 if shared.state.interrupted:

                     break

 

+                if clip_grad:

+                    clip_grad_sched.step(embedding.step)

+            

                 with devices.autocast():

-                    # 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)

-                    loss = shared.sd_model(x, c)[0] / gradient_step

+

+                    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()

@@ -348,6 +389,10 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
                 # go back until we reach gradient accumulation steps

                 if (j + 1) % gradient_step != 0:

                     continue

+                

+                if clip_grad:

+                    clip_grad(embedding.vec, clip_grad_sched.learn_rate)

+

                 scaler.step(optimizer)

                 scaler.update()

                 embedding.step += 1

@@ -366,9 +411,7 @@ def train_embedding(embedding_name, learn_rate, batch_size, gradient_step, data_
                     # 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, checkpoint, embedding_name_every, last_saved_file, remove_cached_checksum=True)

+                    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, {

@@ -458,7 +501,7 @@ Last saved image: {html.escape(last_saved_image)}<br/>
 </p>

 """

         filename = os.path.join(shared.cmd_opts.embeddings_dir, f'{embedding_name}.pt')

-        save_embedding(embedding, checkpoint, embedding_name, filename, remove_cached_checksum=True)

+        save_embedding(embedding, optimizer, checkpoint, embedding_name, filename, remove_cached_checksum=True)

     except Exception:

         print(traceback.format_exc(), file=sys.stderr)

         pass

@@ -470,7 +513,7 @@ Last saved image: {html.escape(last_saved_image)}<br/>
 

     return embedding, filename

 

-def save_embedding(embedding, checkpoint, embedding_name, filename, remove_cached_checksum=True):

+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

@@ -481,6 +524,7 @@ def save_embedding(embedding, checkpoint, embedding_name, filename, remove_cache
         if remove_cached_checksum:

             embedding.cached_checksum = None

         embedding.name = embedding_name

+        embedding.optimizer_state_dict = optimizer.state_dict()

         embedding.save(filename)

     except:

         embedding.sd_checkpoint = old_sd_checkpoint