Add support for --config parameter to kohya_gui.py

kohya_gui.py

@@ -40,7 +40,7 @@ def UI(**kwargs):
         css=css, title=f"Kohya_ss GUI {release}", theme=gr.themes.Default()
     )
     
-    config = KohyaSSGUIConfig()
+    config = KohyaSSGUIConfig(config_file_path=kwargs.get("config_file_path"))
 
     with interface:
         with gr.Tab("Dreambooth"):
@@ -105,6 +105,12 @@ def UI(**kwargs):
 if __name__ == "__main__":
     # torch.cuda.set_per_process_memory_fraction(0.48)
     parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--config",
+        type=str,
+        default="./config.toml",
+        help="Path to the toml config file for interface defaults",
+    )
     parser.add_argument(
         "--listen",
         type=str,
@@ -138,6 +144,7 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     UI(
+        config_file_path=args.config,
         username=args.username,
         password=args.password,
         inbrowser=args.inbrowser,

kohya_gui/class_gui_config.py

@@ -10,13 +10,13 @@ class KohyaSSGUIConfig:
     A class to handle the configuration for the Kohya SS GUI.
     """
 
-    def __init__(self):
+    def __init__(self, config_file_path: str = "./config.toml"):
         """
         Initialize the KohyaSSGUIConfig class.
         """
-        self.config = self.load_config()
+        self.config = self.load_config(config_file_path=config_file_path)
 
-    def load_config(self) -> dict:
+    def load_config(self, config_file_path: str = "./config.toml") -> dict:
         """
         Loads the Kohya SS GUI configuration from a TOML file.
 
@@ -25,16 +25,16 @@ class KohyaSSGUIConfig:
         """
         try:
             # Attempt to load the TOML configuration file from the specified directory.
-            config = toml.load(f"{scriptdir}/config.toml")
-            log.debug(f"Loaded configuration from {scriptdir}/config.toml")
+            config = toml.load(f"{config_file_path}")
+            log.debug(f"Loaded configuration from {config_file_path}")
         except FileNotFoundError:
             # If the config file is not found, initialize `config` as an empty dictionary to handle missing configurations gracefully.
             config = {}
-            log.debug(f"No configuration file found at {scriptdir}/config.toml. Initializing empty configuration.")
+            log.debug(f"No configuration file found at {config_file_path}. Initializing empty configuration.")
 
         return config
 
-    def save_config(self, config: dict):
+    def save_config(self, config: dict, config_file_path: str = "./config.toml"):
         """
         Saves the Kohya SS GUI configuration to a TOML file.
 
@@ -42,7 +42,7 @@ class KohyaSSGUIConfig:
         - config (dict): The configuration data to save.
         """
         # Write the configuration data to the TOML file
-        with open(f"{scriptdir}/config.toml", "w") as f:
+        with open(f"{config_file_path}", "w") as f:
             toml.dump(config, f)
 
     def get(self, key: str, default=None):