Merge pull request #2359 from bmaltais/dev

v24.0.6

.release

@@ -1 +1 @@
-v24.0.5
+v24.0.6

README.md

@@ -42,10 +42,12 @@ The GUI allows you to set the training parameters and generate and run the requi
   - [SDXL training](#sdxl-training)
   - [Masked loss](#masked-loss)
   - [Change History](#change-history)
+    - [2024/04/220 (v24.0.6)](#202404220-v2406)
     - [2024/04/19 (v24.0.5)](#20240419-v2405)
+      - [New Contributors](#new-contributors)
     - [2024/04/18 (v24.0.4)](#20240418-v2404)
       - [What's Changed](#whats-changed)
-      - [New Contributors](#new-contributors)
+      - [New Contributors](#new-contributors-1)
     - [2024/04/24 (v24.0.3)](#20240424-v2403)
     - [2024/04/24 (v24.0.2)](#20240424-v2402)
     - [2024/04/17 (v24.0.1)](#20240417-v2401)
@@ -412,9 +414,20 @@ ControlNet dataset is used to specify the mask. The mask images should be the RG
 
 ## Change History
 
+### 2024/04/220 (v24.0.6)
+
+- Make start and stop buttons visible in headless
+- Add validation for lr and optimizer arguments
+
 ### 2024/04/19 (v24.0.5)
 
-- fdds
+- Hide tensorboard button if tensorflow module is not installed by @bmaltais in <https://github.com/bmaltais/kohya_ss/pull/2347>
+- wd14 captioning issue with undesired tags nor tag replacement by @bmaltais in <https://github.com/bmaltais/kohya_ss/pull/2350>
+- Changed logger checkbox to dropdown, renamed use_wandb -> log_with by @ccharest93 in <https://github.com/bmaltais/kohya_ss/pull/2352>
+
+#### New Contributors
+
+- @ccharest93 made their first contribution in <https://github.com/bmaltais/kohya_ss/pull/2352>
 
 ### 2024/04/18 (v24.0.4)
 
@@ -468,6 +481,7 @@ The `gui.bat` and `gui.sh` scripts now include the `--do_not_use_shell` argument
 #### Miscellaneous
 
 - Made various other minor improvements and bug fixes to enhance overall functionality and user experience.
+- Fixed an issue with existing LoRA network weights were not properly loaded prior to training
 
 ### 2024/04/10 (v23.1.5)
 

assets/style.css

@@ -37,11 +37,11 @@
 #myTensorButton {
     background: radial-gradient(ellipse, #3a99ff, #52c8ff);
     color: white;
-    border: none;
+    border: #296eb8;
 }
 
 #myTensorButtonStop {
     background: radial-gradient(ellipse, #52c8ff, #3a99ff);
     color: black;
-    border: none;
+    border: #296eb8;
 }

config example.toml

@@ -113,7 +113,7 @@ save_state_on_train_end = false           # Save state on train end
 scale_v_pred_loss_like_noise_pred = false # Scale v pred loss like noise pred
 shuffle_caption = false                   # Shuffle captions
 state_dir = "./outputs"                   # Resume from saved training state
-use_wandb = false                         # Use wandb
+log_with = ""                             # Logger to use ["wandb", "tensorboard", "all", ""]
 vae_batch_size = 0                        # VAE batch size
 vae_dir = "./models/vae"                  # VAEs folder path
 v_pred_like_loss = 0                      # V pred like loss weight

kohya_gui/class_command_executor.py

@@ -14,12 +14,19 @@ class CommandExecutor:
     A class to execute and manage commands.
     """
 
-    def __init__(self):
+    def __init__(self, headless: bool = False):
         """
         Initialize the CommandExecutor.
         """
+        self.headless = headless
         self.process = None
-        self.run_state = gr.Textbox(value="", visible=False)
+
+        with gr.Row():
+            self.button_run = gr.Button("Start training", variant="primary")
+
+            self.button_stop_training = gr.Button(
+                "Stop training", visible=self.process is not None or headless, variant="stop"
+            )
 
     def execute_command(self, run_cmd: str, use_shell: bool = False, **kwargs):
         """
@@ -64,16 +71,17 @@ def kill_command(self):
                 # General exception handling for any other errors
                 log.info(f"Error when terminating process: {e}")
         else:
+            self.process = None
             log.info("There is no running process to kill.")
 
-        return gr.Button(visible=True), gr.Button(visible=False)
+        return gr.Button(visible=True), gr.Button(visible=False or self.headless)
 
     def wait_for_training_to_end(self):
         while self.is_running():
             time.sleep(1)
             log.debug("Waiting for training to end...")
         log.info("Training has ended.")
-        return gr.Button(visible=True), gr.Button(visible=False)
+        return gr.Button(visible=True), gr.Button(visible=False or self.headless)
 
     def is_running(self):
         """
@@ -82,4 +90,4 @@ def is_running(self):
         Returns:
         - bool: True if the command is running, False otherwise.
         """
-        return self.process and self.process.poll() is None
+        return self.process is not None and self.process.poll() is None

kohya_gui/class_tensorboard.py

@@ -7,10 +7,11 @@
 from threading import Thread, Event
 from .custom_logging import setup_logging
 
+
 class TensorboardManager:
     DEFAULT_TENSORBOARD_PORT = 6006
 
-    def __init__(self, logging_dir, headless=True, wait_time=5):
+    def __init__(self, logging_dir, headless: bool = False, wait_time=5):
         self.logging_dir = logging_dir
         self.headless = headless
         self.wait_time = wait_time
@@ -25,9 +26,17 @@ def __init__(self, logging_dir, headless=True, wait_time=5):
         self.gradio_interface()
 
     def get_button_states(self, started=False):
-        return gr.Button(visible=not started), gr.Button(visible=started)
+        return gr.Button(visible=not started or self.headless), gr.Button(
+            visible=started or self.headless
+        )
 
     def start_tensorboard(self, logging_dir=None):
+        if self.tensorboard_proc is not None:
+            self.log.info(
+                "Tensorboard is already running. Terminating existing process before starting new one..."
+            )
+            self.stop_tensorboard()
+
         if not os.path.exists(logging_dir) or not os.listdir(logging_dir):
             self.log.error(
                 "Error: logging folder does not exist or does not contain logs."
@@ -46,11 +55,6 @@ def start_tensorboard(self, logging_dir=None):
         ]
 
         self.log.info(run_cmd)
-        if self.tensorboard_proc is not None:
-            self.log.info(
-                "Tensorboard is already running. Terminating existing process before starting new one..."
-            )
-            self.stop_tensorboard()
 
         self.log.info("Starting TensorBoard on port {}".format(self.tensorboard_port))
         try:
@@ -73,7 +77,7 @@ def open_tensorboard_url():
             self.thread = Thread(target=open_tensorboard_url)
             self.thread.start()
 
-        return self.get_button_states(started=True)
+        return self.get_button_states(started=True or self.headless)
 
     def stop_tensorboard(self):
         if self.tensorboard_proc is not None:
@@ -84,34 +88,38 @@ def stop_tensorboard(self):
                 self.log.info("...process stopped")
             except Exception as e:
                 self.log.error("Failed to stop Tensorboard:", e)
-        
+
         if self.thread is not None:
             self.stop_event.set()
             self.thread.join()  # Wait for the thread to finish
             self.thread = None
             self.log.info("Thread terminated successfully.")
 
-        return self.get_button_states(started=False)
+        return self.get_button_states(started=False or self.headless)
 
     def gradio_interface(self):
         try:
-            os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
-            
-            import tensorflow # Attempt to import tensorflow to check if it is installed
-            
+            os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
+
+            import tensorflow  # Attempt to import tensorflow to check if it is installed
+
             visibility = True
-            
+
         except ImportError:
-            self.log.error("tensorflow is not installed, hiding the tensorboard button...")
+            self.log.error(
+                "tensorflow is not installed, hiding the tensorboard button..."
+            )
             visibility = False
-            
+
         with gr.Row():
             button_start_tensorboard = gr.Button(
-                value="Start tensorboard", elem_id="myTensorButton", visible=visibility
+                value="Start tensorboard",
+                elem_id="myTensorButton",
+                visible=visibility or self.headless,
             )
             button_stop_tensorboard = gr.Button(
                 value="Stop tensorboard",
-                visible=False,
+                visible=False or self.headless,
                 elem_id="myTensorButtonStop",
             )
             button_start_tensorboard.click(
@@ -124,4 +132,4 @@ def gradio_interface(self):
                 self.stop_tensorboard,
                 outputs=[button_start_tensorboard, button_stop_tensorboard],
                 show_progress=False,
-            )
+            )

kohya_gui/common_gui.py

@@ -429,7 +429,15 @@ def update_my_data(my_data):
                 pass
 
         my_data.pop(key, None)
-
+
+
+    # Replace the lora_network_weights key with network_weights keeping the original value
+    for key in ["lora_network_weights"]:
+        value = my_data.get(key) # Get original value
+        if value is not None: # Check if the key exists in the dictionary
+            my_data["network_weights"] = value
+            my_data.pop(key, None)
+
     return my_data
 
 
@@ -1490,4 +1498,17 @@ def print_command_and_toml(run_cmd, tmpfilename):
         log.info(toml_file.read())
     log.info(f"end of toml config file: {tmpfilename}")
 
-    save_to_file(command_to_run)
+    save_to_file(command_to_run)
+
+def validate_args_setting(input_string):
+    # Regex pattern to handle multiple conditions:
+    # - Empty string is valid
+    # - Single or multiple key/value pairs with exactly one space between pairs
+    # - No spaces around '=' and no spaces within keys or values
+    pattern = r'^(\S+=\S+)( \S+=\S+)*$|^$'
+    if re.match(pattern, input_string):
+        return True
+    else:
+        log.info(f"'{input_string}' is not a valid settings string.")
+        log.info("A valid settings string must consist of one or more key/value pairs formatted as key=value, with no spaces around the equals sign or within the value. Multiple pairs should be separated by a space.")
+        return False