Add support for resume in Textual Inversion

1 files changed, 54 insertions, 65 deletions

diff --git a/textual_inversion.py b/textual_inversion.py
index a9c3326..11babd8 100644
--- a/textual_inversion.py
+++ b/textual_inversion.py
@@ -170,7 +170,7 @@ def parse_args():
     parser.add_argument(
         "--lr_scheduler",
         type=str,
-        default="one_cycle",
+        default="constant_with_warmup",
         help=(
             'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
             ' "constant", "constant_with_warmup", "one_cycle"]'
@@ -231,14 +231,14 @@ def parse_args():
     parser.add_argument(
         "--checkpoint_frequency",
         type=int,
-        default=500,
-        help="How often to save a checkpoint and sample image",
+        default=5,
+        help="How often to save a checkpoint and sample image (in epochs)",
     )
     parser.add_argument(
         "--sample_frequency",
         type=int,
-        default=100,
-        help="How often to save a checkpoint and sample image",
+        default=1,
+        help="How often to save a checkpoint and sample image (in epochs)",
     )
     parser.add_argument(
         "--sample_image_size",
@@ -294,10 +294,9 @@ def parse_args():
         help="Path to a directory to resume training from (ie, logs/token_name/2022-09-22T23-36-27)"
     )
     parser.add_argument(
-        "--resume_checkpoint",
-        type=str,
-        default=None,
-        help="Path to a specific checkpoint to resume training from (ie, logs/token_name/2022-09-22T23-36-27/checkpoints/something.bin)."
+        "--global_step",
+        type=int,
+        default=0,
     )
     parser.add_argument(
         "--config",
@@ -512,19 +511,10 @@ def main():
     if len(args.placeholder_token) != 0:
         instance_identifier = instance_identifier.format(args.placeholder_token[0])
 
-    global_step_offset = 0
-    if args.resume_from is not None:
-        basepath = Path(args.resume_from)
-        print("Resuming state from %s" % args.resume_from)
-        with open(basepath.joinpath("resume.json"), 'r') as f:
-            state = json.load(f)
-        global_step_offset = state["args"].get("global_step", 0)
-
-        print("We've trained %d steps so far" % global_step_offset)
-    else:
-        now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
-        basepath = Path(args.output_dir).joinpath(slugify(instance_identifier), now)
-        basepath.mkdir(parents=True, exist_ok=True)
+    global_step_offset = args.global_step
+    now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
+    basepath = Path(args.output_dir).joinpath(slugify(instance_identifier), now)
+    basepath.mkdir(parents=True, exist_ok=True)
 
     accelerator = Accelerator(
         log_with=LoggerType.TENSORBOARD,
@@ -557,6 +547,7 @@ def main():
     set_use_memory_efficient_attention_xformers(vae, True)
 
     if args.gradient_checkpointing:
+        unet.enable_gradient_checkpointing()
         text_encoder.gradient_checkpointing_enable()
 
     print(f"Adding text embeddings: {args.placeholder_token}")
@@ -577,14 +568,25 @@ def main():
 
     # Initialise the newly added placeholder token with the embeddings of the initializer token
     token_embeds = text_encoder.get_input_embeddings().weight.data
-    original_token_embeds = token_embeds.detach().clone().to(accelerator.device)
 
-    if args.resume_checkpoint is not None:
-        token_embeds[placeholder_token_id] = torch.load(args.resume_checkpoint)[args.placeholder_token]
-    else:
-        initializer_token_embeddings = text_encoder.get_input_embeddings()(initializer_token_ids)
-        for (token_id, embeddings) in zip(placeholder_token_id, initializer_token_embeddings):
-            token_embeds[token_id] = embeddings
+    if args.resume_from:
+        resumepath = Path(args.resume_from).joinpath("checkpoints")
+
+        for (token_id, token) in zip(placeholder_token_id, args.placeholder_token):
+            embedding_file = resumepath.joinpath(f"{token}_{args.global_step}_end.bin")
+            embedding_data = torch.load(embedding_file, map_location="cpu")
+
+            emb = next(iter(embedding_data.values()))
+            if len(emb.shape) == 1:
+                emb = emb.unsqueeze(0)
+
+            token_embeds[token_id] = emb
+
+    original_token_embeds = token_embeds.clone().to(accelerator.device)
+
+    initializer_token_embeddings = text_encoder.get_input_embeddings()(initializer_token_ids)
+    for (token_id, embeddings) in zip(placeholder_token_id, initializer_token_embeddings):
+        token_embeds[token_id] = embeddings
 
     index_fixed_tokens = torch.arange(len(tokenizer))
     index_fixed_tokens = index_fixed_tokens[~torch.isin(index_fixed_tokens, torch.tensor(placeholder_token_id))]
@@ -891,21 +893,16 @@ def main():
 
                     accelerator.backward(loss)
 
-                    # Keep the token embeddings fixed except the newly added
-                    # embeddings for the concept, as we only want to optimize the concept embeddings
-                    if accelerator.num_processes > 1:
-                        token_embeds = text_encoder.module.get_input_embeddings().weight
-                    else:
-                        token_embeds = text_encoder.get_input_embeddings().weight
-
-                    # Get the index for tokens that we want to freeze
-                    token_embeds.data[index_fixed_tokens, :] = original_token_embeds[index_fixed_tokens, :]
-
                     optimizer.step()
                     if not accelerator.optimizer_step_was_skipped:
                         lr_scheduler.step()
                     optimizer.zero_grad(set_to_none=True)
 
+                    # Let's make sure we don't update any embedding weights besides the newly added token
+                    with torch.no_grad():
+                        text_encoder.get_input_embeddings(
+                        ).weight[index_fixed_tokens] = original_token_embeds[index_fixed_tokens]
+
                     loss = loss.detach().item()
                     train_loss += loss
 
@@ -916,19 +913,6 @@ def main():
 
                     global_step += 1
 
-                    if global_step % args.sample_frequency == 0:
-                        sample_checkpoint = True
-
-                    if global_step % args.checkpoint_frequency == 0 and global_step > 0 and accelerator.is_main_process:
-                        local_progress_bar.clear()
-                        global_progress_bar.clear()
-
-                        checkpointer.checkpoint(global_step + global_step_offset, "training")
-                        save_args(basepath, args, {
-                            "global_step": global_step + global_step_offset,
-                            "resume_checkpoint": f"{basepath}/checkpoints/last.bin"
-                        })
-
                 logs = {"train/loss": loss, "lr": lr_scheduler.get_last_lr()[0]}
 
                 accelerator.log(logs, step=global_step)
@@ -992,24 +976,30 @@ def main():
             local_progress_bar.clear()
             global_progress_bar.clear()
 
-            if min_val_loss > val_loss:
-                accelerator.print(
-                    f"Global step {global_step}: Validation loss reached new minimum: {min_val_loss:.2e} -> {val_loss:.2e}")
-                checkpointer.checkpoint(global_step + global_step_offset, "milestone")
-                min_val_loss = val_loss
+            if accelerator.is_main_process:
+                if min_val_loss > val_loss:
+                    accelerator.print(
+                        f"Global step {global_step}: Validation loss reached new minimum: {min_val_loss:.2e} -> {val_loss:.2e}")
+                    checkpointer.checkpoint(global_step + global_step_offset, "milestone")
+                    min_val_loss = val_loss
+
+                if epoch % args.checkpoint_frequency == 0:
+                    checkpointer.checkpoint(global_step + global_step_offset, "training")
+                    save_args(basepath, args, {
+                        "global_step": global_step + global_step_offset
+                    })
 
-            if sample_checkpoint and accelerator.is_main_process:
-                checkpointer.save_samples(
-                    global_step + global_step_offset,
-                    args.resolution, args.resolution, 7.5, 0.0, args.sample_steps)
+                if epoch % args.sample_frequency == 0:
+                    checkpointer.save_samples(
+                        global_step + global_step_offset,
+                        args.resolution, args.resolution, 7.5, 0.0, args.sample_steps)
 
         # Create the pipeline using using the trained modules and save it.
         if accelerator.is_main_process:
             print("Finished! Saving final checkpoint and resume state.")
             checkpointer.checkpoint(global_step + global_step_offset, "end")
             save_args(basepath, args, {
-                "global_step": global_step + global_step_offset,
-                "resume_checkpoint": f"{basepath}/checkpoints/last.bin"
+                "global_step": global_step + global_step_offset
             })
             accelerator.end_training()
 
@@ -1018,8 +1008,7 @@ def main():
             print("Interrupted, saving checkpoint and resume state...")
             checkpointer.checkpoint(global_step + global_step_offset, "end")
             save_args(basepath, args, {
-                "global_step": global_step + global_step_offset,
-                "resume_checkpoint": f"{basepath}/checkpoints/last.bin"
+                "global_step": global_step + global_step_offset
             })
             accelerator.end_training()
         quit()