forked from bmaltais/kohya_ss
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathclass_gui_config.py
93 lines (76 loc) · 3.13 KB
/
class_gui_config.py
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
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