Merge pull request #2450 from bmaltais/dev

v24.1.0
bmaltais · May 5, 2024 · 95f8883 · 95f8883
2 parents 415d7b2 + d26dad2
commit 95f8883
Show file tree

Hide file tree

Showing 39 changed files with 1,060 additions and 943 deletions.
diff --git a/.github/workflows/typos.yaml b/.github/workflows/typos.yaml
@@ -18,4 +18,4 @@ jobs:
       - uses: actions/checkout@v4
 
       - name: typos-action
-        uses: crate-ci/typos@v1.19.0
+        uses: crate-ci/typos@v1.21.0
diff --git a/.release b/.release
@@ -1 +1 @@
-v24.0.9
+v24.1.0
diff --git a/README.md b/README.md
diff --git a/kohya_gui.py b/kohya_gui.py
@@ -78,13 +78,11 @@ def UI(**kwargs):
                 reg_data_dir_input=reg_data_dir_input,
                 output_dir_input=output_dir_input,
                 logging_dir_input=logging_dir_input,
-                enable_copy_info_button=True,
                 headless=headless,
                 config=config,
-                use_shell_flag=use_shell_flag,
             )
             with gr.Tab("LoRA"):
-                _ = LoRATools(headless=headless, use_shell_flag=use_shell_flag)
+                _ = LoRATools(headless=headless)
         with gr.Tab("About"):
             gr.Markdown(f"kohya_ss GUI release {release}")
             with gr.Tab("README"):

diff --git a/kohya_gui/basic_caption_gui.py b/kohya_gui/basic_caption_gui.py
@@ -1,5 +1,4 @@
 import gradio as gr
-from easygui import msgbox
 import subprocess
 from .common_gui import (
     get_folder_path,
@@ -28,7 +27,6 @@ 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.
@@ -46,72 +44,78 @@ def caption_images(
     Returns:
         None
     """
-    # Check if images_dir is provided
+    # Check if images_dir and caption_ext are provided
+    missing_parameters = []
     if not images_dir:
-        msgbox(
-            "Image folder is missing. Please provide the directory containing the images to caption."
-        )
-        return
-
-    # Check if caption_ext is provided
+        missing_parameters.append("image directory")
     if not caption_ext:
-        msgbox("Please provide an extension for the caption files.")
+        missing_parameters.append("caption file extension")
+
+    if missing_parameters:
+        log.info(
+            "The following parameter(s) are missing: {}. "
+            "Please provide these to proceed with captioning the images.".format(", ".join(missing_parameters))
+        )
         return
 
     # Log the captioning process
     if caption_text:
         log.info(f"Captioning files in {images_dir} with {caption_text}...")
 
         # Build the command to run caption.py
-        run_cmd = rf'"{PYTHON}" "{scriptdir}/tools/caption.py"'
-        run_cmd += f' --caption_text="{caption_text}"'
+        run_cmd = [
+            rf"{PYTHON}",
+            rf"{scriptdir}/tools/caption.py",
+            "--caption_text",
+            caption_text,
+        ]
 
         # Add optional flags to the command
         if overwrite:
-            run_cmd += f" --overwrite"
+            run_cmd.append("--overwrite")
         if caption_ext:
-            run_cmd += f' --caption_file_ext="{caption_ext}"'
+            run_cmd.append("--caption_file_ext")
+            run_cmd.append(caption_ext)
 
-        run_cmd += f' "{images_dir}"'
+        run_cmd.append(rf"{images_dir}")
 
-        # Log the command
-        log.info(run_cmd)
+        # Reconstruct the safe command string for display
+        command_to_run = " ".join(run_cmd)
+        log.info(f"Executing command: {command_to_run}")
 
         # Set the environment variable for the Python path
         env = os.environ.copy()
         env["PYTHONPATH"] = (
-            f"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
+            rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
         )
         env["TF_ENABLE_ONEDNN_OPTS"] = "0"
 
-        log.info(f"Executing command: {run_cmd} with shell={use_shell}")
-
         # Run the command in the sd-scripts folder context
-        subprocess.run(run_cmd, env=env, shell=use_shell)
-
+        subprocess.run(run_cmd, env=env, shell=False)
 
     # Check if overwrite option is enabled
     if overwrite:
-        # Add prefix and postfix to caption files
-        if prefix or postfix:
+        # Add prefix and postfix to caption files or find and replace text in caption files
+        if prefix or postfix or find_text:
+            # Add prefix and/or postfix to caption files
             add_pre_postfix(
                 folder=images_dir,
                 caption_file_ext=caption_ext,
                 prefix=prefix,
                 postfix=postfix,
             )
-        # Find and replace text in caption files
-        if find_text:
-            find_replace(
-                folder_path=images_dir,
-                caption_file_ext=caption_ext,
-                search_text=find_text,
-                replace_text=replace_text,
-            )
+            # Replace specified text in caption files if find and replace text is provided
+            if find_text and replace_text:
+                find_replace(
+                    folder_path=images_dir,
+                    caption_file_ext=caption_ext,
+                    search_text=find_text,
+                    replace_text=replace_text,
+                )
     else:
         # Show a message if modification is not possible without overwrite option enabled
         if prefix or postfix:
-            msgbox(
+            log.info(
                 'Could not modify caption files with requested change because the "Overwrite existing captions in folder" option is not selected.'
             )
 
@@ -120,7 +124,7 @@ def caption_images(
 
 
 # Gradio UI
-def gradio_basic_caption_gui_tab(headless=False, default_images_dir=None, use_shell: bool = False):
+def gradio_basic_caption_gui_tab(headless=False, default_images_dir=None):
     """
     Creates a Gradio tab for basic image captioning.
 
@@ -200,6 +204,7 @@ def list_images_dirs(path):
                 choices=[".cap", ".caption", ".txt"],
                 value=".txt",
                 interactive=True,
+                allow_custom_value=True,
             )
             # Checkbox to overwrite existing captions
             overwrite = gr.Checkbox(
@@ -258,7 +263,6 @@ def list_images_dirs(path):
                     postfix,
                     find_text,
                     replace_text,
-                    gr.Checkbox(value=use_shell, visible=False),
                 ],
                 show_progress=False,
             )

diff --git a/kohya_gui/blip_caption_gui.py b/kohya_gui/blip_caption_gui.py
@@ -1,5 +1,4 @@
 import gradio as gr
-from easygui import msgbox
 import subprocess
 import os
 import sys
@@ -23,7 +22,6 @@ def caption_images(
     beam_search: bool,
     prefix: str = "",
     postfix: str = "",
-    use_shell: bool = False,
 ) -> None:
     """
     Automatically generates captions for images in the specified directory using the BLIP model.
@@ -47,59 +45,60 @@ def caption_images(
     """
     # Check if the image folder is provided
     if not train_data_dir:
-        msgbox("Image folder is missing...")
+        log.info("Image folder is missing...")
         return
 
     # Check if the caption file extension is provided
     if not caption_file_ext:
-        msgbox("Please provide an extension for the caption files.")
+        log.info("Please provide an extension for the caption files.")
         return
 
     log.info(f"Captioning files in {train_data_dir}...")
 
     # Construct the command to run make_captions.py
-    run_cmd = [fr'"{PYTHON}"', fr'"{scriptdir}/sd-scripts/finetune/make_captions.py"']
+    run_cmd = [rf"{PYTHON}", rf"{scriptdir}/sd-scripts/finetune/make_captions.py"]
 
     # Add required arguments
-    run_cmd.append('--batch_size')
+    run_cmd.append("--batch_size")
     run_cmd.append(str(batch_size))
-    run_cmd.append('--num_beams')
+    run_cmd.append("--num_beams")
     run_cmd.append(str(num_beams))
-    run_cmd.append('--top_p')
+    run_cmd.append("--top_p")
     run_cmd.append(str(top_p))
-    run_cmd.append('--max_length')
+    run_cmd.append("--max_length")
     run_cmd.append(str(max_length))
-    run_cmd.append('--min_length')
+    run_cmd.append("--min_length")
     run_cmd.append(str(min_length))
 
     # Add optional flags to the command
     if beam_search:
         run_cmd.append("--beam_search")
     if caption_file_ext:
-        run_cmd.append('--caption_extension')
+        run_cmd.append("--caption_extension")
         run_cmd.append(caption_file_ext)
 
     # Add the directory containing the training data
-    run_cmd.append(fr'"{train_data_dir}"')
+    run_cmd.append(rf"{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")
+    run_cmd.append("--caption_weights")
+    run_cmd.append(
+        rf"https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_large_caption.pth"
+    )
 
     # Set up the environment
     env = os.environ.copy()
     env["PYTHONPATH"] = (
-        f"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
+        rf"{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(command_to_run, env=env, shell=use_shell, cwd=f"{scriptdir}/sd-scripts")
+    log.info(f"Executing command: {command_to_run}")
 
+    # Run the command in the sd-scripts folder context
+    subprocess.run(run_cmd, env=env, shell=False, cwd=rf"{scriptdir}/sd-scripts")
 
     # Add prefix and postfix
     add_pre_postfix(
@@ -117,7 +116,7 @@ def caption_images(
 ###
 
 
-def gradio_blip_caption_gui_tab(headless=False, default_train_dir=None, use_shell: bool = False):
+def gradio_blip_caption_gui_tab(headless=False, default_train_dir=None):
     from .common_gui import create_refresh_button
 
     default_train_dir = (
@@ -167,6 +166,7 @@ def list_train_dirs(path):
                 choices=[".cap", ".caption", ".txt"],
                 value=".txt",
                 interactive=True,
+                allow_custom_value=True,
             )
 
             prefix = gr.Textbox(
@@ -207,7 +207,6 @@ def list_train_dirs(path):
                 beam_search,
                 prefix,
                 postfix,
-                gr.Checkbox(value=use_shell, visible=False),
             ],
             show_progress=False,
         )

diff --git a/kohya_gui/class_command_executor.py b/kohya_gui/class_command_executor.py
@@ -28,7 +28,7 @@ def __init__(self, headless: bool = False):
                 "Stop training", visible=self.process is not None or headless, variant="stop"
             )
 
-    def execute_command(self, run_cmd: str, use_shell: bool = False, **kwargs):
+    def execute_command(self, run_cmd: str, **kwargs):
         """
         Execute a command if no other command is currently running.
 
@@ -44,10 +44,10 @@ def execute_command(self, run_cmd: str, use_shell: bool = False, **kwargs):
 
             # 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}")
+            log.info(f"Executing command: {command_to_run}")
 
             # Execute the command securely
-            self.process = subprocess.Popen(command_to_run, **kwargs, shell=use_shell)
+            self.process = subprocess.Popen(run_cmd, **kwargs)
             log.info("Command executed.")
 
     def kill_command(self):

diff --git a/kohya_gui/class_lora_tab.py b/kohya_gui/class_lora_tab.py
@@ -14,15 +14,14 @@ class LoRATools:
     def __init__(
         self,
         headless: bool = False,
-        use_shell_flag: bool = False,
     ):
         gr.Markdown("This section provide various LoRA tools...")
-        gradio_extract_dylora_tab(headless=headless, use_shell=use_shell_flag)
-        gradio_convert_lcm_tab(headless=headless, use_shell=use_shell_flag)
-        gradio_extract_lora_tab(headless=headless, use_shell=use_shell_flag)
-        gradio_extract_lycoris_locon_tab(headless=headless, use_shell=use_shell_flag)
-        gradio_merge_lora_tab = GradioMergeLoRaTab(use_shell=use_shell_flag)
-        gradio_merge_lycoris_tab(headless=headless, use_shell=use_shell_flag)
-        gradio_svd_merge_lora_tab(headless=headless, use_shell=use_shell_flag)
+        gradio_extract_dylora_tab(headless=headless)
+        gradio_convert_lcm_tab(headless=headless)
+        gradio_extract_lora_tab(headless=headless)
+        gradio_extract_lycoris_locon_tab(headless=headless)
+        gradio_merge_lora_tab = GradioMergeLoRaTab()
+        gradio_merge_lycoris_tab(headless=headless)
+        gradio_svd_merge_lora_tab(headless=headless)
         gradio_resize_lora_tab(headless=headless)
         gradio_verify_lora_tab(headless=headless)