Update training info display

CongLeSolutionX · Apr 17, 2024 · 43087f7 · 43087f7
1 parent 8bda4f2
commit 43087f7
Show file tree

Hide file tree

Showing 4 changed files with 56 additions and 25 deletions.
diff --git a/kohya_gui/dreambooth_gui.py b/kohya_gui/dreambooth_gui.py
@@ -530,8 +530,11 @@ def train_model(
                 lr_warmup_steps = 0
         else:
             lr_warmup_steps = 0
-
-        max_train_steps_info = f"Max train steps: {max_train_steps}"
+
+        if max_train_steps == 0:
+            max_train_steps_info = f"Max train steps: 0. sd-scripts will therefore default to 1600. Please specify a different value if required."
+        else:
+            max_train_steps_info = f"Max train steps: {max_train_steps}"
     else:
         if train_data_dir == "":
             log.error("Train data dir is empty")
@@ -604,19 +607,23 @@ def train_model(
             )
             max_train_steps_info = f"max_train_steps ({total_steps} / {train_batch_size} / {gradient_accumulation_steps} * {epoch} * {reg_factor}) = {max_train_steps}"
         else:
-            max_train_steps_info = f"Max train steps: {max_train_steps}"
+            if max_train_steps == 0:
+                max_train_steps_info = f"Max train steps: 0. sd-scripts will therefore default to 1600. Please specify a different value if required."
+            else:
+                max_train_steps_info = f"Max train steps: {max_train_steps}"
 
         if lr_warmup != 0:
             lr_warmup_steps = round(float(int(lr_warmup) * int(max_train_steps) / 100))
         else:
             lr_warmup_steps = 0
 
         log.info(f"Total steps: {total_steps}")
-        log.info(f"Train batch size: {train_batch_size}")
-        log.info(f"Gradient accumulation steps: {gradient_accumulation_steps}")
-        log.info(f"Epoch: {epoch}")
-        log.info(max_train_steps_info)
-        log.info(f"lr_warmup_steps = {lr_warmup_steps}")
+
+    log.info(f"Train batch size: {train_batch_size}")
+    log.info(f"Gradient accumulation steps: {gradient_accumulation_steps}")
+    log.info(f"Epoch: {epoch}")
+    log.info(max_train_steps_info)
+    log.info(f"lr_warmup_steps = {lr_warmup_steps}")
 
     run_cmd = [fr'"{get_executable_path("accelerate")}"', "launch"]
 

diff --git a/kohya_gui/finetune_gui.py b/kohya_gui/finetune_gui.py
@@ -559,6 +559,11 @@ def train_model(
         log.info(
             "Dataset config toml file used, skipping caption json file, image buckets, total_steps, train_batch_size, gradient_accumulation_steps, epoch, reg_factor, max_train_steps creation..."
         )
+
+        if max_train_steps == 0:
+            max_train_steps_info = f"Max train steps: 0. sd-scripts will therefore default to 1600. Please specify a different value if required."
+        else:
+            max_train_steps_info = f"Max train steps: {max_train_steps}"
     else:
         # create caption json file
         if generate_caption_database:
@@ -675,8 +680,13 @@ def train_model(
         # Divide by two because flip augmentation create two copied of the source images
         if flip_aug and max_train_steps:
             max_train_steps = int(math.ceil(float(max_train_steps) / 2))
+
+        if max_train_steps == 0:
+            max_train_steps_info = f"Max train steps: 0. sd-scripts will therefore default to 1600. Please specify a different value if required."
+        else:
+            max_train_steps_info = f"Max train steps: {max_train_steps}"
 
-    log.info(f"max_train_steps = {max_train_steps}")
+    log.info(max_train_steps_info)
 
     if max_train_steps != 0:
         lr_warmup_steps = round(float(int(lr_warmup) * int(max_train_steps) / 100))

diff --git a/kohya_gui/lora_gui.py b/kohya_gui/lora_gui.py
@@ -777,7 +777,10 @@ def train_model(
             stop_text_encoder_training = 0
             lr_warmup_steps = 0
 
-        max_train_steps_info = f"Max train steps: {max_train_steps}"
+        if max_train_steps == 0:
+            max_train_steps_info = f"Max train steps: 0. sd-scripts will therefore default to 1600. Please specify a different value if required."
+        else:
+            max_train_steps_info = f"Max train steps: {max_train_steps}"
 
     else:
         if train_data_dir == "":
@@ -851,7 +854,10 @@ def train_model(
             )
             max_train_steps_info = f"max_train_steps ({total_steps} / {train_batch_size} / {gradient_accumulation_steps} * {epoch} * {reg_factor}) = {max_train_steps}"
         else:
-            max_train_steps_info = f"Max train steps: {max_train_steps}"
+            if max_train_steps == 0:
+                max_train_steps_info = f"Max train steps: 0. sd-scripts will therefore default to 1600. Please specify a different value if required."
+            else:
+                max_train_steps_info = f"Max train steps: {max_train_steps}"
 
         # calculate stop encoder training
         if stop_text_encoder_training_pct == 0:
@@ -867,12 +873,13 @@ def train_model(
             lr_warmup_steps = 0
 
         log.info(f"Total steps: {total_steps}")
-        log.info(f"Train batch size: {train_batch_size}")
-        log.info(f"Gradient accumulation steps: {gradient_accumulation_steps}")
-        log.info(f"Epoch: {epoch}")
-        log.info(max_train_steps_info)
-        log.info(f"stop_text_encoder_training = {stop_text_encoder_training}")
-        log.info(f"lr_warmup_steps = {lr_warmup_steps}")
+
+    log.info(f"Train batch size: {train_batch_size}")
+    log.info(f"Gradient accumulation steps: {gradient_accumulation_steps}")
+    log.info(f"Epoch: {epoch}")
+    log.info(max_train_steps_info)
+    log.info(f"stop_text_encoder_training = {stop_text_encoder_training}")
+    log.info(f"lr_warmup_steps = {lr_warmup_steps}")
 
     run_cmd = [rf'"{get_executable_path("accelerate")}"', "launch"]
 

diff --git a/kohya_gui/textual_inversion_gui.py b/kohya_gui/textual_inversion_gui.py
@@ -545,7 +545,10 @@ def train_model(
             stop_text_encoder_training = 0
             lr_warmup_steps = 0
 
-        max_train_steps_info = f"Max train steps: {max_train_steps}"
+        if max_train_steps == 0:
+            max_train_steps_info = f"Max train steps: 0. sd-scripts will therefore default to 1600. Please specify a different value if required."
+        else:
+            max_train_steps_info = f"Max train steps: {max_train_steps}"
 
     else:
         if train_data_dir == "":
@@ -619,7 +622,10 @@ def train_model(
             )
             max_train_steps_info = f"max_train_steps ({total_steps} / {train_batch_size} / {gradient_accumulation_steps} * {epoch} * {reg_factor}) = {max_train_steps}"
         else:
-            max_train_steps_info = f"Max train steps: {max_train_steps}"
+            if max_train_steps == 0:
+                max_train_steps_info = f"Max train steps: 0. sd-scripts will therefore default to 1600. Please specify a different value if required."
+            else:
+                max_train_steps_info = f"Max train steps: {max_train_steps}"
 
         # calculate stop encoder training
         if stop_text_encoder_training_pct == 0:
@@ -635,12 +641,13 @@ def train_model(
             lr_warmup_steps = 0
 
         log.info(f"Total steps: {total_steps}")
-        log.info(f"Train batch size: {train_batch_size}")
-        log.info(f"Gradient accumulation steps: {gradient_accumulation_steps}")
-        log.info(f"Epoch: {epoch}")
-        log.info(max_train_steps_info)
-        log.info(f"stop_text_encoder_training = {stop_text_encoder_training}")
-        log.info(f"lr_warmup_steps = {lr_warmup_steps}")
+
+    log.info(f"Train batch size: {train_batch_size}")
+    log.info(f"Gradient accumulation steps: {gradient_accumulation_steps}")
+    log.info(f"Epoch: {epoch}")
+    log.info(max_train_steps_info)
+    log.info(f"stop_text_encoder_training = {stop_text_encoder_training}")
+    log.info(f"lr_warmup_steps = {lr_warmup_steps}")
 
     run_cmd = [fr'"{get_executable_path("accelerate")}"', "launch"]