Add support for user managed path config

Author: bmaltais

.gitignore

@@ -51,3 +51,5 @@ outputs
 dataset/**
 !dataset/**/
 !dataset/**/.gitkeep
+# models
+# data

README.md

@@ -373,6 +373,7 @@ The documentation in this section will be moved to a separate document later.
 - Add support for `wandb_run_name`, `log_tracker_name` and `log_tracker_config` parameters under the advanced section.
 - Update sd-scripts to v0.8.5
 - Improve code
+- Add support for custom path defaults. Simply copy the `config example.toml` file found in the root of the repo to `config.toml` and edit the different values to your taste.
 
 ### 2024/03/13 (v23.0.11)
 

config example.toml

@@ -0,0 +1,13 @@
+# Copy this file and name it config.toml
+# Edit the values to suit your needs
+
+# Default folders location
+models_dir = "./models"    # Pretrained model name or path
+train_data_dir = "./data"   # Image folder (containing training images subfolders) / Image folder (containing training images)
+output_dir = "./outputs"    # Output directory for trained model
+reg_data_dir = "./data/reg" # Regularisation directory
+logging_dir = "./logs"      # Logging directory
+config_dir = "./presets"    # Load/Save Config file
+
+# Example custom folder location
+# models_dir = "e:/models"    # Pretrained model name or path

kohya_gui/class_configuration_file.py

@@ -1,7 +1,11 @@
 import gradio as gr
 import os
+import toml
+from .common_gui import list_files, scriptdir, create_refresh_button, load_kohya_ss_gui_config
+from .custom_logging import setup_logging
 
-from .common_gui import list_files, scriptdir, create_refresh_button
+# Set up logging
+log = setup_logging()
 
 
 class ConfigurationFile:
@@ -19,11 +23,11 @@ def __init__(self, headless: bool = False, config_dir: str = None):
         """
 
         self.headless = headless
+        
+        config = load_kohya_ss_gui_config()
 
         # Sets the directory for storing configuration files, defaults to a 'presets' folder within the script directory.
-        self.current_config_dir = (
-            config_dir if config_dir is not None else os.path.join(scriptdir, "presets")
-        )
+        self.current_config_dir = config.get('config_dir', os.path.join(scriptdir, "presets"))
 
         # Initialize the GUI components for configuration.
         self.create_config_gui()
@@ -38,7 +42,7 @@ def list_config_dir(self, path: str) -> list:
         Returns:
         - list: A list of directories.
         """
-        self.current_config_dir = path
+        self.current_config_dir = path if not path == "" else "."
         # Lists all .json files in the current configuration directory, used for populating dropdown choices.
         return list(list_files(self.current_config_dir, exts=[".json"], all=True))
 

kohya_gui/class_folders.py

@@ -1,31 +1,29 @@
 import gradio as gr
 import os
-from .common_gui import get_folder_path, scriptdir, list_dirs, create_refresh_button
+from .common_gui import get_folder_path, scriptdir, list_dirs, create_refresh_button, load_kohya_ss_gui_config
 
 class Folders:
     """
     A class to handle folder operations in the GUI.
     """
-    def __init__(self, finetune: bool = False, data_dir: str = None, output_dir: str = None, logging_dir: str = None, reg_data_dir: str = None, headless: bool = False):
+    def __init__(self, finetune: bool = False, headless: bool = False):
         """
         Initialize the Folders class.
 
         Parameters:
         - finetune (bool): Whether to finetune the model.
-        - data_dir (str): The directory for data.
-        - output_dir (str): The directory for output.
-        - logging_dir (str): The directory for logging.
-        - reg_data_dir (str): The directory for regularization data.
         - headless (bool): Whether to run in headless mode.
         """
         self.headless = headless
         self.finetune = finetune
+        
+        # Load kohya_ss GUI configs from config.toml if it exist
+        config = load_kohya_ss_gui_config()
 
         # Set default directories if not provided
-        self.current_data_dir = data_dir if data_dir is not None else os.path.join(scriptdir, "data")
-        self.current_output_dir = output_dir if output_dir is not None else os.path.join(scriptdir, "outputs")
-        self.current_logging_dir = logging_dir if logging_dir is not None else os.path.join(scriptdir, "logs")
-        self.current_reg_data_dir = reg_data_dir if reg_data_dir is not None else os.path.join(scriptdir, "reg")
+        self.current_output_dir = config.get('output_dir', os.path.join(scriptdir, "outputs"))
+        self.current_logging_dir = config.get('logging_dir', os.path.join(scriptdir, "logs"))
+        self.current_reg_data_dir = config.get('reg_data_dir', os.path.join(scriptdir, "reg"))
 
         # Create directories if they don't exist
         self.create_directory_if_not_exists(self.current_output_dir)
@@ -44,18 +42,6 @@ def create_directory_if_not_exists(self, directory: str) -> None:
         if directory is not None and directory.strip() != "" and not os.path.exists(directory):
             os.makedirs(directory, exist_ok=True)
 
-    def list_data_dirs(self, path: str) -> list:
-        """
-        List directories in the data directory.
-
-        Parameters:
-        - path (str): The path to list directories from.
-
-        Returns:
-        - list: A list of directories.
-        """
-        self.current_data_dir = path
-        return list(list_dirs(path))
 
     def list_output_dirs(self, path: str) -> list:
         """
@@ -67,7 +53,7 @@ def list_output_dirs(self, path: str) -> list:
         Returns:
         - list: A list of directories.
         """
-        self.current_output_dir = path
+        self.current_output_dir = path if not path == "" else "."
         return list(list_dirs(path))
 
     def list_logging_dirs(self, path: str) -> list:
@@ -80,7 +66,7 @@ def list_logging_dirs(self, path: str) -> list:
         Returns:
         - list: A list of directories.
         """
-        self.current_logging_dir = path
+        self.current_logging_dir = path if not path == "" else "."
         return list(list_dirs(path))
 
     def list_reg_data_dirs(self, path: str) -> list:
@@ -93,7 +79,7 @@ def list_reg_data_dirs(self, path: str) -> list:
         Returns:
         - list: A list of directories.
         """
-        self.current_reg_data_dir = path
+        self.current_reg_data_dir = path if not path == "" else "."
         return list(list_dirs(path))
 
     def create_folders_gui(self) -> None:
@@ -131,7 +117,7 @@ def create_folders_gui(self) -> None:
                 allow_custom_value=True,
             )
             # Refresh button for regularisation directory
-            create_refresh_button(self.reg_data_dir, lambda: None, lambda: {"choices": [""] + self.list_data_dirs(self.current_data_dir)}, "open_folder_small")
+            create_refresh_button(self.reg_data_dir, lambda: None, lambda: {"choices": [""] + self.list_reg_data_dirs(self.current_reg_data_dir)}, "open_folder_small")
             # Regularisation directory button
             self.reg_data_dir_folder = gr.Button(
                 '📂', elem_id='open_folder_small', elem_classes=["tool"], visible=(not self.headless)
@@ -173,7 +159,7 @@ def create_folders_gui(self) -> None:
             )
             # Change event for regularisation directory dropdown
             self.reg_data_dir.change(
-                fn=lambda path: gr.Dropdown(choices=[""] + self.list_data_dirs(path)),
+                fn=lambda path: gr.Dropdown(choices=[""] + self.list_reg_data_dirs(path)),
                 inputs=self.reg_data_dir,
                 outputs=self.reg_data_dir,
                 show_progress=False,