File size: 3,296 Bytes
fb83c5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
import toml
from .common_gui import scriptdir
from .custom_logging import setup_logging
# Set up logging
log = setup_logging()
class KohyaSSGUIConfig:
"""
A class to handle the configuration for the Kohya SS GUI.
"""
def __init__(self, config_file_path: str = "./config.toml"):
"""
Initialize the KohyaSSGUIConfig class.
"""
self.config = self.load_config(config_file_path=config_file_path)
def load_config(self, config_file_path: str = "./config.toml") -> dict:
"""
Loads the Kohya SS GUI configuration from a TOML file.
Returns:
dict: The configuration data loaded from the TOML file.
"""
try:
# Attempt to load the TOML configuration file from the specified directory.
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 {config_file_path}. Initializing empty configuration."
)
return config
def save_config(self, config: dict, config_file_path: str = "./config.toml"):
"""
Saves the Kohya SS GUI configuration to a TOML file.
Parameters:
- config (dict): The configuration data to save.
"""
# Write the configuration data to the TOML file
with open(f"{config_file_path}", "w", encoding="utf-8") as f:
toml.dump(config, f)
def get(self, key: str, default=None):
"""
Retrieves the value of a specified key from the configuration data.
Parameters:
- key (str): The key to retrieve the value for.
- default: The default value to return if the key is not found.
Returns:
The value associated with the key, or the default value if the key is not found.
"""
# Split the key into a list of keys if it contains a dot (.)
keys = key.split(".")
# Initialize `data` with the entire configuration data
data = self.config
# Iterate over the keys to access nested values
for k in keys:
log.debug(k)
# If the key is not found in the current data, return the default value
if k not in data:
log.debug(
f"Key '{key}' not found in configuration. Returning default value."
)
return default
# Update `data` to the value associated with the current key
data = data.get(k)
# Return the final value
log.debug(f"Returned {data}")
return data
def is_config_loaded(self) -> bool:
"""
Checks if the configuration was loaded from a file.
Returns:
bool: True if the configuration was loaded from a file, False otherwise.
"""
is_loaded = self.config != {}
log.debug(f"Configuration was loaded from file: {is_loaded}")
return is_loaded
|