Move config for gui to seperate class

Add support for more config paths

Author: bmaltais

.gitignore

@@ -41,15 +41,12 @@ test/ft
 # Temporary requirements
 requirements_tmp_for_setup.txt
 
-# Version specific
-0.13.3
-
 *.npz
 presets/*/user_presets/*
 inputs
 outputs
 dataset/**
 !dataset/**/
 !dataset/**/.gitkeep
-# models
-# data
+models
+data

config example.toml

@@ -2,12 +2,15 @@
 # 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
+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
+log_tracker_config_dir = "./logs" # Log tracker configs directory
+state_dir = "./outputs"           # Resume from saved training state
+vae_dir = "./models/vae"          # VAEs folder path
 
 # Example custom folder location
-# models_dir = "e:/models"    # Pretrained model name or path
+# models_dir = "e:/models"    # Pretrained model name or path

config.toml

@@ -0,0 +1,16 @@
+# 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
+log_tracker_config_dir = "./logs" # Log tracker configs directory
+state_dir = "e:/models"           # Resume from saved training state
+vae_dir = "./models/vae"          # VAEs folder path
+
+# Example custom folder location
+# models_dir = "e:/models"    # Pretrained model name or path

kohya_gui.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import os
 import argparse
+from kohya_gui.class_gui_config import KohyaSSGUIConfig
 from dreambooth_gui import dreambooth_tab
 from finetune_gui import finetune_tab
 from textual_inversion_gui import ti_tab
@@ -24,7 +25,7 @@ def UI(**kwargs):
 
     if os.path.exists("./style.css"):
         with open(os.path.join("./style.css"), "r", encoding="utf8") as file:
-            log.info("Load CSS...")
+            log.debug("Load CSS...")
             css += file.read() + "\n"
 
     if os.path.exists("./.release"):
@@ -38,6 +39,8 @@ def UI(**kwargs):
     interface = gr.Blocks(
         css=css, title=f"Kohya_ss GUI {release}", theme=gr.themes.Default()
     )
+    
+    config = KohyaSSGUIConfig()
 
     with interface:
         with gr.Tab("Dreambooth"):
@@ -46,13 +49,13 @@ def UI(**kwargs):
                 reg_data_dir_input,
                 output_dir_input,
                 logging_dir_input,
-            ) = dreambooth_tab(headless=headless)
+            ) = dreambooth_tab(headless=headless, config=config)
         with gr.Tab("LoRA"):
-            lora_tab(headless=headless)
+            lora_tab(headless=headless, config=config)
         with gr.Tab("Textual Inversion"):
-            ti_tab(headless=headless)
+            ti_tab(headless=headless, config=config)
         with gr.Tab("Finetuning"):
-            finetune_tab(headless=headless)
+            finetune_tab(headless=headless, config=config)
         with gr.Tab("Utilities"):
             utilities_tab(
                 train_data_dir_input=train_data_dir_input,

kohya_gui/class_advanced_training.py

@@ -30,8 +30,7 @@ def __init__(
         headless: bool = False,
         finetuning: bool = False,
         training_type: str = "",
-        default_vae_dir: str = "",
-        default_output_dir: str = "",
+        config:dict = {},
     ) -> None:
         """
         Initializes the AdvancedTraining class with given settings.
@@ -46,16 +45,12 @@ def __init__(
         self.headless = headless
         self.finetuning = finetuning
         self.training_type = training_type
+        self.config = config
 
         # Determine the current directories for VAE and output, falling back to defaults if not specified.
-        current_vae_dir = (
-            default_vae_dir if default_vae_dir else os.path.join(scriptdir, "vae")
-        )
-        current_state_dir = (
-            default_output_dir
-            if default_output_dir
-            else os.path.join(scriptdir, "outputs")
-        )
+        self.current_vae_dir = self.config.get("vae_dir", "./models/vae")
+        self.current_state_dir = self.config.get("state_dir", "./outputs")
+        self.current_log_tracker_config_dir = self.config.get("log_tracker_config_dir", "./logs")
 
         # Define the behavior for changing noise offset type.
         def noise_offset_type_change(
@@ -95,21 +90,20 @@ def noise_offset_type_change(
             self.prior_loss_weight = gr.Number(label="Prior loss weight", value=1.0)
 
             def list_vae_files(path):
-                nonlocal current_vae_dir
-                current_vae_dir = path
+                self.current_vae_dir = path if not path == "" else "."
                 return list(list_files(path, exts=[".ckpt", ".safetensors"], all=True))
 
             self.vae = gr.Dropdown(
-                label="VAE (Optional. path to checkpoint of vae to replace for training)",
+                label="VAE (Optional: Path to checkpoint of vae for training)",
                 interactive=True,
-                choices=[""] + list_vae_files(current_vae_dir),
+                choices=[""] + list_vae_files(self.current_vae_dir),
                 value="",
                 allow_custom_value=True,
             )
             create_refresh_button(
                 self.vae,
                 lambda: None,
-                lambda: {"choices": list_vae_files(current_vae_dir)},
+                lambda: {"choices": [""] + list_vae_files(self.current_vae_dir)},
                 "open_folder_small",
             )
             self.vae_button = gr.Button(
@@ -222,11 +216,6 @@ def full_options_update(full_fp16, full_bf16):
                 label="Memory efficient attention", value=False
             )
         with gr.Row():
-            # This use_8bit_adam element should be removed in a future release as it is no longer used
-            # use_8bit_adam = gr.Checkbox(
-            #     label='Use 8bit adam', value=False, visible=False
-            # )
-            # self.xformers = gr.Checkbox(label='Use xformers', value=True, info='Use xformers for CrossAttention')
             self.xformers = gr.Dropdown(
                 label="CrossAttention",
                 choices=["none", "sdpa", "xformers"],
@@ -348,21 +337,20 @@ def full_options_update(full_fp16, full_bf16):
             self.save_state = gr.Checkbox(label="Save training state", value=False)
 
             def list_state_dirs(path):
-                nonlocal current_state_dir
-                current_state_dir = path
+                self.current_state_dir = path if not path == "" else "."
                 return list(list_dirs(path))
 
             self.resume = gr.Dropdown(
                 label='Resume from saved training state (path to "last-state" state folder)',
-                choices=[""] + list_state_dirs(current_state_dir),
+                choices=[""] + list_state_dirs(self.current_state_dir),
                 value="",
                 interactive=True,
                 allow_custom_value=True,
             )
             create_refresh_button(
                 self.resume,
                 lambda: None,
-                lambda: {"choices": list_state_dirs(current_state_dir)},
+                lambda: {"choices": [""] + list_state_dirs(self.current_state_dir)},
                 "open_folder_small",
             )
             self.resume_button = gr.Button(
@@ -418,6 +406,10 @@ def list_state_dirs(path):
                 info="The name of the specific wandb session",
             )
         with gr.Group(), gr.Row():
+            def list_log_tracker_config_files(path):
+                self.current_log_tracker_config_dir = path if not path == "" else "."
+                return list(list_files(path, exts=[".json"], all=True))
+    
             self.log_tracker_name = gr.Textbox(
                 label="Log tracker name",
                 value="",
@@ -426,7 +418,7 @@ def list_state_dirs(path):
             )
             self.log_tracker_config = gr.Dropdown(
                 label="Log tracker config",
-                choices=[""] + list_state_dirs(current_state_dir),
+                choices=[""] + list_log_tracker_config_files(self.current_log_tracker_config_dir),
                 value="",
                 info="Path to tracker config file to use for logging",
                 interactive=True,
@@ -435,7 +427,7 @@ def list_state_dirs(path):
             create_refresh_button(
                 self.log_tracker_config,
                 lambda: None,
-                lambda: {"choices": list_state_dirs(current_state_dir)},
+                lambda: {"choices": [""] + list_log_tracker_config_files(self.current_log_tracker_config_dir)},
                 "open_folder_small",
             )
             self.log_tracker_config_button = gr.Button(
@@ -447,7 +439,7 @@ def list_state_dirs(path):
                 show_progress=False,
             )
             self.log_tracker_config.change(
-                fn=lambda path: gr.Dropdown(choices=[""] + list_state_dirs(path)),
+                fn=lambda path: gr.Dropdown(choices=[""] + list_log_tracker_config_files(path)),
                 inputs=self.log_tracker_config,
                 outputs=self.log_tracker_config,
                 show_progress=False,

kohya_gui/class_configuration_file.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import os
-import toml
-from .common_gui import list_files, scriptdir, create_refresh_button, load_kohya_ss_gui_config
+from .common_gui import list_files, scriptdir, create_refresh_button
 from .custom_logging import setup_logging
 
 # Set up logging
@@ -13,7 +12,7 @@ class ConfigurationFile:
     A class to handle configuration file operations in the GUI.
     """
 
-    def __init__(self, headless: bool = False, config_dir: str = None):
+    def __init__(self, headless: bool = False, config_dir: str = None, config:dict = {}):
         """
         Initialize the ConfigurationFile class.
 
@@ -24,10 +23,10 @@ def __init__(self, headless: bool = False, config_dir: str = None):
 
         self.headless = headless
 
-        config = load_kohya_ss_gui_config()
+        self.config = config
 
         # Sets the directory for storing configuration files, defaults to a 'presets' folder within the script directory.
-        self.current_config_dir = config.get('config_dir', os.path.join(scriptdir, "presets"))
+        self.current_config_dir = self.config.get('config_dir', os.path.join(scriptdir, "presets"))
 
         # Initialize the GUI components for configuration.
         self.create_config_gui()