1 files changed, 7 insertions, 5 deletions

diff --git a/train_lora.py b/train_lora.py
index 073e939..b8c7396 100644
--- a/train_lora.py
+++ b/train_lora.py
@@ -945,6 +945,8 @@ def main():
     if accelerator.is_main_process:
         accelerator.init_trackers(lora_project)
 
+    lora_sample_output_dir = output_dir / lora_project / "samples"
+
     while True:
         if training_iter >= args.auto_cycles:
             response = input("Run another cycle? [y/n] ")
@@ -995,8 +997,7 @@ def main():
             train_epochs=num_train_epochs,
         )
 
-        lora_checkpoint_output_dir = output_dir / lora_project / f"{training_iter + 1}" / "model"
-        lora_sample_output_dir = output_dir / lora_project / f"{training_iter + 1}" / "samples"
+        lora_checkpoint_output_dir = output_dir / lora_project / f"model_{training_iter + 1}"
 
         trainer(
             strategy=lora_strategy,
@@ -1007,6 +1008,7 @@ def main():
             num_train_epochs=num_train_epochs,
             gradient_accumulation_steps=args.gradient_accumulation_steps,
             global_step_offset=training_iter * num_train_steps,
+            initial_samples=training_iter == 0,
             # --
             group_labels=group_labels,
             sample_output_dir=lora_sample_output_dir,
@@ -1015,11 +1017,11 @@ def main():
         )
 
         training_iter += 1
-        if args.learning_rate_emb is not None:
+        if learning_rate_emb is not None:
             learning_rate_emb *= args.cycle_decay
-        if args.learning_rate_unet is not None:
+        if learning_rate_unet is not None:
             learning_rate_unet *= args.cycle_decay
-        if args.learning_rate_text is not None:
+        if learning_rate_text is not None:
             learning_rate_text *= args.cycle_decay
 
     accelerator.end_training()