mirror of https://github.com/bmaltais/kohya_ss
78 lines
2.6 KiB
Python
78 lines
2.6 KiB
Python
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):
|
|
"""
|
|
Initialize the KohyaSSGUIConfig class.
|
|
"""
|
|
self.config = self.load_config()
|
|
|
|
def load_config(self) -> 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"{scriptdir}/config.toml")
|
|
log.debug(f"Loaded configuration from {scriptdir}/config.toml")
|
|
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.")
|
|
|
|
return config
|
|
|
|
def save_config(self, config: dict):
|
|
"""
|
|
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"{scriptdir}/config.toml", "w") 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
|