Reverting all changes for max_train_steps

Author: bmaltais

kohya_gui/dreambooth_gui.py

@@ -654,8 +654,6 @@ def train_model(
         gr.Button(visible=False or headless),
         gr.Textbox(value=train_state_value),
     ]
-    
-    max_train_steps_info = "Automatic by sd-scripts"
 
     if executor.is_running():
         log.error("Training is already running. Can't start another training session.")
@@ -783,23 +781,23 @@ def train_model(
 
         log.info(f"Regularization factor: {reg_factor}")
 
-        # if max_train_steps == 0:
-        #     # calculate max_train_steps
-        #     max_train_steps = int(
-        #         math.ceil(
-        #             float(total_steps)
-        #             / int(train_batch_size)
-        #             / int(gradient_accumulation_steps)
-        #             * int(epoch)
-        #             * int(reg_factor)
-        #         )
-        #     )
-        #     max_train_steps_info = f"max_train_steps ({total_steps} / {train_batch_size} / {gradient_accumulation_steps} * {epoch} * {reg_factor}) = {max_train_steps}"
-        # else:
-        #     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 max_train_steps == 0:
+            # calculate max_train_steps
+            max_train_steps = int(
+                math.ceil(
+                    float(total_steps)
+                    / int(train_batch_size)
+                    / int(gradient_accumulation_steps)
+                    * int(epoch)
+                    * int(reg_factor)
+                )
+            )
+            max_train_steps_info = f"max_train_steps ({total_steps} / {train_batch_size} / {gradient_accumulation_steps} * {epoch} * {reg_factor}) = {max_train_steps}"
+        else:
+            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"Total steps: {total_steps}")
 
@@ -1470,4 +1468,4 @@ def dreambooth_tab(
             folders.reg_data_dir,
             folders.output_dir,
             folders.logging_dir,
-        )
+        )

kohya_gui/finetune_gui.py

@@ -685,8 +685,6 @@ def train_model(
         gr.Button(visible=False or headless),
         gr.Textbox(value=train_state_value),
     ]
-    
-    max_train_steps_info = "Automatic by sd-scripts"
 
     if executor.is_running():
         log.error("Training is already running. Can't start another training session.")
@@ -848,16 +846,16 @@ def train_model(
         repeats = int(image_num) * int(dataset_repeats)
         log.info(f"repeats = {str(repeats)}")
 
-        # if max_train_steps == 0:
-        #     # calculate max_train_steps
-        #     max_train_steps = int(
-        #         math.ceil(
-        #             float(repeats)
-        #             / int(train_batch_size)
-        #             / int(gradient_accumulation_steps)
-        #             * int(epoch)
-        #         )
-        #     )
+        if max_train_steps == 0:
+            # calculate max_train_steps
+            max_train_steps = int(
+                math.ceil(
+                    float(repeats)
+                    / int(train_batch_size)
+                    / int(gradient_accumulation_steps)
+                    * int(epoch)
+                )
+            )
 
         # Divide by two because flip augmentation create two copied of the source images
         if flip_aug and max_train_steps:
@@ -1634,4 +1632,4 @@ def list_presets(path):
         if os.path.exists(top_level_path):
             with open(os.path.join(top_level_path), "r", encoding="utf-8") as file:
                 guides_top_level = file.read() + "\n"
-            gr.Markdown(guides_top_level)
+            gr.Markdown(guides_top_level)

kohya_gui/lora_gui.py

@@ -883,8 +883,6 @@ def train_model(
         gr.Button(visible=False or headless),
         gr.Textbox(value=train_state_value),
     ]
-    
-    max_train_steps_info = "Automatic by sd-scripts"
 
     if executor.is_running():
         log.error("Training is already running. Can't start another training session.")
@@ -1078,7 +1076,7 @@ def train_model(
 
         log.info(f"Regularization factor: {reg_factor}")
 
-        if (max_train_steps == 0) and (stop_text_encoder_training != 0):
+        if max_train_steps == 0:
             # calculate max_train_steps
             max_train_steps = int(
                 math.ceil(
@@ -1096,9 +1094,13 @@ def train_model(
             else:
                 max_train_steps_info = f"Max train steps: {max_train_steps}"
 
-        stop_text_encoder_training = math.ceil(
-            float(max_train_steps) / 100 * int(stop_text_encoder_training)
-        ) if stop_text_encoder_training != 0 else 0
+        # calculate stop encoder training
+        if stop_text_encoder_training == 0:
+            stop_text_encoder_training = 0
+        else:
+            stop_text_encoder_training = math.ceil(
+                float(max_train_steps) / 100 * int(stop_text_encoder_training)
+            )
 
     # Calculate lr_warmup_steps
     if lr_warmup_steps > 0:
@@ -2855,4 +2857,4 @@ def update_LoRA_settings(
         folders.reg_data_dir,
         folders.output_dir,
         folders.logging_dir,
-    )
+    )

kohya_gui/textual_inversion_gui.py

@@ -521,8 +521,6 @@ def train_model(
         gr.Textbox(value=train_state_value),
     ]
 
-    max_train_steps_info = "Automatic by sd-scripts"
-    
     if executor.is_running():
         log.error("Training is already running. Can't start another training session.")
         return TRAIN_BUTTON_VISIBLE
@@ -666,9 +664,22 @@ def train_model(
         log.info(f"Regularization factor: {reg_factor}")
 
         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."
+            # calculate max_train_steps
+            max_train_steps = int(
+                math.ceil(
+                    float(total_steps)
+                    / int(train_batch_size)
+                    / int(gradient_accumulation_steps)
+                    * int(epoch)
+                    * int(reg_factor)
+                )
+            )
+            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:
@@ -1065,6 +1076,10 @@ def list_embedding_files(path):
                             step=1,
                             label="Vectors",
                         )
+                        # max_train_steps = gr.Textbox(
+                        #     label='Max train steps',
+                        #     placeholder='(Optional) Maximum number of steps',
+                        # )
                         template = gr.Dropdown(
                             label="Template",
                             choices=[
@@ -1294,4 +1309,4 @@ def list_embedding_files(path):
             folders.reg_data_dir,
             folders.output_dir,
             folders.logging_dir,
-        )
+        )