Implement LR Warmup disabling when constant scheduler is used

README.md

@@ -410,6 +410,7 @@ ControlNet dataset is used to specify the mask. The mask images should be the RG
 - Added a HuggingFace section to all trainers tabs, enabling users to authenticate and utilize HuggingFace's powerful AI models.
 - Converted the Graphical User Interface (GUI) to use the configuration TOML file format to pass arguments to sd-scripts. This change improves security by eliminating the need for sensitive information to be passed through the command-line interface (CLI).
 - Made various other minor improvements and bug fixes to enhance the overall functionality and user experience.
+- Disabled LR Warmup when using the Constant LR Scheduler to prevent traceback errors with sd-scripts.
 
 ### 2024/04/10 (v23.1.5)
 

kohya_gui/class_basic_training.py

@@ -1,6 +1,9 @@
 import gradio as gr
 from typing import Tuple
+from .custom_logging import setup_logging
 
+# Set up logging
+log = setup_logging()
 
 class BasicTraining:
     """
@@ -44,6 +47,7 @@ def __init__(
         self.finetuning = finetuning
         self.dreambooth = dreambooth
         self.config = config
+        self.old_lr_warmup = 0
 
         # Initialize the UI components
         self.initialize_ui_components()
@@ -162,6 +166,9 @@ def init_lr_and_optimizer_controls(self) -> None:
                 ],
                 value=self.config.get("basic.lr_scheduler", self.lr_scheduler_value),
             )
+
+
+
             # Initialize the optimizer dropdown
             self.optimizer = gr.Dropdown(
                 label="Optimizer",
@@ -278,6 +285,26 @@ def init_learning_rate_controls(self) -> None:
                 maximum=100,
                 step=1,
             )
+
+            def lr_scheduler_changed(scheduler, value):
+                if scheduler == "constant":
+                    self.old_lr_warmup = value
+                    value = 0
+                    interactive=False
+                    info="Can't use LR warmup with LR Scheduler constant... setting to 0 and disabling field..."
+                else:
+                    if self.old_lr_warmup != 0:
+                        value = self.old_lr_warmup
+                        self.old_lr_warmup = 0
+                    interactive=True
+                    info=""
+                return gr.Slider(value=value, interactive=interactive, info=info)
+
+            self.lr_scheduler.change(
+                lr_scheduler_changed,
+                inputs=[self.lr_scheduler, self.lr_warmup],
+                outputs=self.lr_warmup,
+            )
 
     def init_scheduler_controls(self) -> None:
         """

kohya_gui/class_command_executor.py

@@ -37,7 +37,7 @@ def execute_command(self, run_cmd: str, **kwargs):
 
             # Reconstruct the safe command string for display
             command_to_run = ' '.join(run_cmd)
-            log.info(f"Executings command: {command_to_run}")
+            log.info(f"Executing command: {command_to_run}")
 
             # Execute the command securely
             self.process = subprocess.Popen(run_cmd, **kwargs)