Improve process execution under Linux (and windows)

README.md

@@ -410,24 +410,26 @@ ControlNet dataset is used to specify the mask. The mask images should be the RG
 
 #### Enhancements
 
-- **User Interface:** Converted the GUI to utilize a configuration TOML file for passing arguments to sd-scripts, enhancing security by avoiding the command-line interface (CLI) for sensitive information.
-- **Training Tools:** Enhanced the functionality of the training and TensorBoard buttons to offer a more intuitive user experience.
-- **HuggingFace Integration:** Added a HuggingFace section to all trainers' tabs, allowing users to authenticate and leverage HuggingFace's advanced AI models.
-- **Gradio Upgrade:** Upgraded to gradio version 4.20.0, resolving a bug affecting the runpod platform.
-- **Metadata Support:** Introduced support for metadata capture within the GUI.
+- **User Interface:** Transitioned the GUI to use a TOML file for argument passing to sd-scripts, significantly enhancing security by eliminating the need for command-line interface (CLI) use for sensitive data.
+- **Training Tools:** Improved the training and TensorBoard buttons to provide a more intuitive user experience.
+- **HuggingFace Integration:** Integrated a HuggingFace section in all trainer tabs, enabling authentication and use of HuggingFace's advanced AI models.
+- **Gradio Upgrade:** Upgraded Gradio to version 4.20.0 to fix a previously identified bug impacting the runpod platform.
+- **Metadata Support:** Added functionality for metadata capture within the GUI.
 
 #### Security and Stability
 
-- **Code Refactoring:** Rewrote significant portions of the code to address security vulnerabilities, including the removal of the `shell=True` parameter from process calls.
-- **Scheduler Update:** Disabled LR Warmup when using the Constant LR Scheduler to avoid traceback errors with sd-scripts.
+- **Code Refactoring:** Extensively rewrote the code to address various security vulnerabilities, including removing the `shell=True` parameter from process calls.
+- **Scheduler Update:** Disabled LR Warmup when using the Constant LR Scheduler to prevent traceback errors associated with sd-scripts.
 
 #### Shell Execution
 
-- **Conditional Shell Usage:** Implemented support for optional shell usage when running external sd-scripts commands to accommodate specific platform requirements, following security enhancements.
-  - **How to Enable Shell Execution:**
-    1. Start the GUI with the `--use_shell` option.
-    2. Set the `use_shell` value in the config.toml file to `true`, allowing the GUI to apply user-specified settings at startup.
-    - **Note:** The `--use_shell` option will take precedence over settings in the config.toml file.
+- **Conditional Shell Usage:** Added support for optional shell usage when executing external sd-scripts commands, tailored to meet specific platform needs and recent security updates.
+
+The `gui.bat` and `gui.sh` scripts now include the `--do_not_use_shell` argument to prevent shell execution (`shell=True`) during external process handling. Unix-like systems automatically set `use_shell` to True internally, as required for proper execution of external commands. To enforce disabling shell execution, use the `--do_not_use_shell` argument.
+
+- **How to Enable Shell Execution via Config File:**
+  1. In the `config.toml` file, set `use_shell` to `true` to enable shell usage as per GUI startup settings.
+  **Note:** The `--do_not_use_shell` option will override the `config.toml` settings, setting `use_shell` to False even if it is set to True in the config file.
 
 #### Miscellaneous
 

kohya_gui.py

@@ -42,9 +42,16 @@ def UI(**kwargs):
     if config.is_config_loaded():
         log.info(f"Loaded default GUI values from '{kwargs.get('config')}'...")
 
-    use_shell_flag = kwargs.get("use_shell", False)
-    if use_shell_flag == False:
-        use_shell_flag = config.get("settings.use_shell", False)
+    use_shell_flag = False
+    if os.name == "posix":
+        use_shell_flag = True
+
+    if config.get("settings.use_shell", False):
+        use_shell_flag = True
+
+    if kwargs.get("do_not_use_shell", False):
+        use_shell_flag = False
+
     if use_shell_flag:
         log.info("Using shell=True when running external commands...")
 
@@ -159,7 +166,7 @@ def UI(**kwargs):
     parser.add_argument("--use-rocm", action="store_true", help="Use ROCm environment")
 
     parser.add_argument(
-        "--use_shell", action="store_true", help="Use shell environment"
+        "--do_not_use_shell", action="store_true", help="Enforce not to use shell=True when running external commands"
     )
 
     parser.add_argument(

kohya_gui/basic_caption_gui.py

@@ -28,6 +28,7 @@ def caption_images(
     postfix: str,
     find_text: str,
     replace_text: str,
+    use_shell: bool = False,
 ):
     """
     Captions images in a given directory with a given caption text.
@@ -62,7 +63,7 @@ def caption_images(
         log.info(f"Captioning files in {images_dir} with {caption_text}...")
 
         # Build the command to run caption.py
-        run_cmd = [PYTHON, f"{scriptdir}/tools/caption.py"]
+        run_cmd = [PYTHON, fr'"{scriptdir}/tools/caption.py"']
 
         # Add required arguments
         run_cmd.append('--caption_text')
@@ -76,10 +77,7 @@ def caption_images(
             run_cmd.append(caption_ext)
 
         # Add the directory containing the images
-        run_cmd.append(images_dir)
-
-        # Log the command
-        log.info(' '.join(run_cmd))
+        run_cmd.append(fr'"{images_dir}"')
 
         # Set the environment variable for the Python path
         env = os.environ.copy()
@@ -88,8 +86,13 @@ def caption_images(
         )
         env["TF_ENABLE_ONEDNN_OPTS"] = "0"
 
-        # Run the command based on the operating system
-        subprocess.run(run_cmd, env=env)
+        # Reconstruct the safe command string for display
+        command_to_run = " ".join(run_cmd)
+        log.info(f"Executing command: {command_to_run} with shell={use_shell}")
+
+        # Run the command in the sd-scripts folder context
+        subprocess.run(command_to_run, env=env, shell=use_shell)
+
 
     # Check if overwrite option is enabled
     if overwrite:
@@ -121,7 +124,7 @@ def caption_images(
 
 
 # Gradio UI
-def gradio_basic_caption_gui_tab(headless=False, default_images_dir=None):
+def gradio_basic_caption_gui_tab(headless=False, default_images_dir=None, use_shell: bool = False):
     """
     Creates a Gradio tab for basic image captioning.
 
@@ -259,6 +262,7 @@ def list_images_dirs(path):
                     postfix,
                     find_text,
                     replace_text,
+                    gr.Checkbox(value=use_shell, visible=False),
                 ],
                 show_progress=False,
             )

kohya_gui/blip_caption_gui.py

@@ -58,7 +58,7 @@ def caption_images(
     log.info(f"Captioning files in {train_data_dir}...")
 
     # Construct the command to run make_captions.py
-    run_cmd = [PYTHON, f"{scriptdir}/sd-scripts/finetune/make_captions.py"]
+    run_cmd = [PYTHON, fr'"{scriptdir}/sd-scripts/finetune/make_captions.py"']
 
     # Add required arguments
     run_cmd.append('--batch_size')
@@ -80,24 +80,25 @@ def caption_images(
         run_cmd.append(caption_file_ext)
 
     # Add the directory containing the training data
-    run_cmd.append(train_data_dir)
+    run_cmd.append(fr'"{train_data_dir}"')
 
     # Add URL for caption model weights
     run_cmd.append('--caption_weights')
     run_cmd.append("https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_large_caption.pth")
 
-    # Log the command
-    log.info(' '.join(run_cmd))
-
     # Set up the environment
     env = os.environ.copy()
     env["PYTHONPATH"] = (
         f"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
     )
     env["TF_ENABLE_ONEDNN_OPTS"] = "0"
 
+    # Reconstruct the safe command string for display
+    command_to_run = " ".join(run_cmd)
+    log.info(f"Executing command: {command_to_run} with shell={use_shell}")
+
     # Run the command in the sd-scripts folder context
-    subprocess.run(run_cmd, env=env, shell=use_shell, cwd=f"{scriptdir}/sd-scripts")
+    subprocess.run(command_to_run, env=env, shell=use_shell, cwd=f"{scriptdir}/sd-scripts")
 
 
     # Add prefix and postfix

kohya_gui/class_command_executor.py

@@ -32,15 +32,15 @@ def execute_command(self, run_cmd: str, use_shell: bool = False, **kwargs):
         if self.process and self.process.poll() is None:
             log.info("The command is already running. Please wait for it to finish.")
         else:
-            # for i, item in enumerate(run_cmd):
-            #     log.info(f"{i}: {item}")
+            for i, item in enumerate(run_cmd):
+                log.info(f"{i}: {item}")
 
             # Reconstruct the safe command string for display
             command_to_run = " ".join(run_cmd)
             log.info(f"Executing command: {command_to_run} with shell={use_shell}")
 
             # Execute the command securely
-            self.process = subprocess.Popen(run_cmd, **kwargs, shell=use_shell)
+            self.process = subprocess.Popen(command_to_run, **kwargs, shell=use_shell)
             log.info("Command executed.")
 
     def kill_command(self):

kohya_gui/convert_lcm_gui.py

@@ -57,7 +57,7 @@ def convert_lcm(
     run_cmd.append("--lora-scale")
     run_cmd.append(str(lora_scale))
     run_cmd.append("--model")
-    run_cmd.append(model_path)
+    run_cmd.append(rf'"{model_path}"')
     run_cmd.append("--name")
     run_cmd.append(name)
 
@@ -67,18 +67,21 @@ def convert_lcm(
     if model_type == "SSD-1B":
         run_cmd.append("--ssd-1b")
 
-    # Log the command
-    log.info(" ".join(run_cmd))
-
     # Set up the environment
     env = os.environ.copy()
     env["PYTHONPATH"] = (
         rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
     )
     env["TF_ENABLE_ONEDNN_OPTS"] = "0"
 
-    # Run the command
-    subprocess.run(run_cmd, env=env, shell=use_shell)
+    # Reconstruct the safe command string for display
+    command_to_run = " ".join(run_cmd)
+    log.info(f"Executing command: {command_to_run} with shell={use_shell}")
+
+    # Run the command in the sd-scripts folder context
+    subprocess.run(
+        command_to_run, env=env, shell=use_shell
+    )
 
     # Return a success message
     log.info("Done extracting...")

kohya_gui/convert_model_gui.py

@@ -50,7 +50,8 @@ def convert_model(
         return
 
     run_cmd = [
-        PYTHON, f"{scriptdir}/sd-scripts/tools/convert_diffusers20_original_sd.py"
+        PYTHON,
+        fr'"{scriptdir}/sd-scripts/tools/convert_diffusers20_original_sd.py"',
     ]
 
     v1_models = [
@@ -70,7 +71,7 @@ def convert_model(
         run_cmd.append(f"--{target_save_precision_type}")
 
     if target_model_type == "diffuser" or target_model_type == "diffuser_safetensors":
-        run_cmd.append('--reference_model')
+        run_cmd.append("--reference_model")
         run_cmd.append(source_model_type)
 
     if target_model_type == "diffuser_safetensors":
@@ -81,7 +82,7 @@ def convert_model(
         run_cmd.append("--unet_use_linear_projection")
 
     # Add the source model input path
-    run_cmd.append(source_model_input)
+    run_cmd.append(fr'"{source_model_input}"')
 
     # Determine the target model path
     if target_model_type == "diffuser" or target_model_type == "diffuser_safetensors":
@@ -95,10 +96,7 @@ def convert_model(
         )
 
     # Add the target model path
-    run_cmd.append(target_model_path)
-
-    # Log the command
-    log.info(' '.join(run_cmd))
+    run_cmd.append(fr'"{target_model_path}"')
 
     env = os.environ.copy()
     env["PYTHONPATH"] = (
@@ -107,9 +105,14 @@ def convert_model(
     # Adding an example of an environment variable that might be relevant
     env["TF_ENABLE_ONEDNN_OPTS"] = "0"
 
-    # Run the command
-    subprocess.run(run_cmd, env=env, shell=use_shell)
+    # Reconstruct the safe command string for display
+    command_to_run = " ".join(run_cmd)
+    log.info(f"Executing command: {command_to_run} with shell={use_shell}")
 
+    # Run the command in the sd-scripts folder context
+    subprocess.run(
+        command_to_run, env=env, shell=use_shell
+    )
 
 
 ###

kohya_gui/dreambooth_gui.py

@@ -628,7 +628,7 @@ def train_model(
         lr_warmup_steps = 0
     log.info(f"lr_warmup_steps = {lr_warmup_steps}")
 
-    run_cmd = [get_executable_path("accelerate"), "launch"]
+    run_cmd = [fr'"{get_executable_path("accelerate")}"', "launch"]
 
     run_cmd = AccelerateLaunch.run_cmd(
         run_cmd=run_cmd,
@@ -647,9 +647,9 @@ def train_model(
     )
 
     if sdxl:
-        run_cmd.append(f"{scriptdir}/sd-scripts/sdxl_train.py")
+        run_cmd.append(fr'"{scriptdir}/sd-scripts/sdxl_train.py"')
     else:
-        run_cmd.append(f"{scriptdir}/sd-scripts/train_db.py")
+        run_cmd.append(fr'"{scriptdir}/sd-scripts/train_db.py"')
 
     if max_data_loader_n_workers == "" or None:
         max_data_loader_n_workers = 0
@@ -798,7 +798,7 @@ def train_model(
             log.error(f"Failed to write TOML file: {toml_file.name}")
 
     run_cmd.append(f"--config_file")
-    run_cmd.append(tmpfilename)
+    run_cmd.append(fr'"{tmpfilename}"')
 
     # Initialize a dictionary with always-included keyword arguments
     kwargs_for_training = {

kohya_gui/extract_lora_from_dylora_gui.py

@@ -52,27 +52,28 @@ def extract_dylora(
 
     run_cmd = [
         PYTHON,
-        f"{scriptdir}/sd-scripts/networks/extract_lora_from_dylora.py",
+        rf'"{scriptdir}/sd-scripts/networks/extract_lora_from_dylora.py"',
         "--save_to",
-        save_to,
+        rf'"{save_to}"',
         "--model",
-        model,
+        rf'"{model}"',
         "--unit",
         str(unit),
     ]
 
-    # Log the command
-    log.info(" ".join(run_cmd))
-
     env = os.environ.copy()
     env["PYTHONPATH"] = (
         rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
     )
     # Example environment variable adjustment for the Python environment
     env["TF_ENABLE_ONEDNN_OPTS"] = "0"
 
-    # Run the command
-    subprocess.run(run_cmd, env=env, shell=use_shell)
+    # Reconstruct the safe command string for display
+    command_to_run = " ".join(run_cmd)
+    log.info(f"Executing command: {command_to_run} with shell={use_shell}")
+
+    # Run the command in the sd-scripts folder context
+    subprocess.run(command_to_run, env=env, shell=use_shell)
 
     log.info("Done extracting DyLoRA...")
 

kohya_gui/extract_lora_gui.py

@@ -75,17 +75,17 @@ def extract_lora(
 
     run_cmd = [
         PYTHON,
-        f"{scriptdir}/sd-scripts/networks/extract_lora_from_models.py",
+        fr'"{scriptdir}/sd-scripts/networks/extract_lora_from_models.py"',
         "--load_precision",
         load_precision,
         "--save_precision",
         save_precision,
         "--save_to",
-        save_to,
+        fr'"{save_to}"',
         "--model_org",
-        model_org,
+        fr'"{model_org}"',
         "--model_tuned",
-        model_tuned,
+        fr'"{model_tuned}"',
         "--dim",
         str(dim),
         "--device",
@@ -110,18 +110,20 @@ def extract_lora(
         run_cmd.append("--load_tuned_model_to")
         run_cmd.append(load_tuned_model_to)
 
-    # Log the command
-    log.info(" ".join(run_cmd))
-
     env = os.environ.copy()
     env["PYTHONPATH"] = (
         rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
     )
     # Adding an example of another potentially relevant environment variable
     env["TF_ENABLE_ONEDNN_OPTS"] = "0"
 
-    # Run the command
-    subprocess.run(run_cmd, env=env, shell=use_shell)
+    # Reconstruct the safe command string for display
+    command_to_run = " ".join(run_cmd)
+    log.info(f"Executing command: {command_to_run} with shell={use_shell}")
+
+    # Run the command in the sd-scripts folder context
+    subprocess.run(command_to_run, env=env, shell=use_shell)
+
 
 
 ###