79 lines
2.2 KiB
Python
79 lines
2.2 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
|
||
|
|
"""Checks and sets default values for config.json before starting the container."""
|
||
|
|
|
||
|
|
import json
|
||
|
|
import re
|
||
|
|
import os.path
|
||
|
|
import sys
|
||
|
|
|
||
|
|
DEFAULT_FILEPATH = '/data/config/auto/config.json'
|
||
|
|
|
||
|
|
DEFAULT_OUTDIRS = {
|
||
|
|
"outdir_samples": "",
|
||
|
|
"outdir_txt2img_samples": "/output/txt2img",
|
||
|
|
"outdir_img2img_samples": "/output/img2img",
|
||
|
|
"outdir_extras_samples": "/output/extras",
|
||
|
|
"outdir_grids": "",
|
||
|
|
"outdir_txt2img_grids": "/output/txt2img-grids",
|
||
|
|
"outdir_img2img_grids": "/output/img2img-grids",
|
||
|
|
"outdir_save": "/output/saved",
|
||
|
|
"outdir_init_images": "/output/init-images",
|
||
|
|
}
|
||
|
|
RE_VALID_OUTDIR = re.compile(r"(^/output(/\.?[\w\-\_]+)+/?$)|(^\s?$)")
|
||
|
|
|
||
|
|
DEFAULT_OTHER = {
|
||
|
|
"font": "DejaVuSans.ttf",
|
||
|
|
}
|
||
|
|
|
||
|
|
def dict_to_json_file(target_file: str, data: dict):
|
||
|
|
"""Write dictionary to specified json file"""
|
||
|
|
|
||
|
|
with open(target_file, 'w') as f:
|
||
|
|
json.dump(data, f)
|
||
|
|
|
||
|
|
def json_file_to_dict(config_file: str) -> dict|None:
|
||
|
|
"""Load json file into a dictionary. Return None if file does not exist."""
|
||
|
|
|
||
|
|
if os.path.isfile(config_file):
|
||
|
|
with open(config_file, 'r') as f:
|
||
|
|
return json.load(f)
|
||
|
|
else:
|
||
|
|
return None
|
||
|
|
|
||
|
|
def replace_if_invalid(value: str, replacement: str, pattern: str|re.Pattern[str]) -> str:
|
||
|
|
"""Returns original value if valid, fallback value if invalid"""
|
||
|
|
|
||
|
|
if re.match(pattern, value):
|
||
|
|
return value
|
||
|
|
else:
|
||
|
|
return replacement
|
||
|
|
|
||
|
|
def check_and_replace_config(config_file: str, target_file: str = None):
|
||
|
|
"""Checks given file for invalid values. Replaces those with fallback values (default: overwrites file)."""
|
||
|
|
|
||
|
|
# Get current user config, or empty if file does not exists
|
||
|
|
data = json_file_to_dict(config_file) or {}
|
||
|
|
|
||
|
|
# Check and fix output directories
|
||
|
|
for k, def_val in DEFAULT_OUTDIRS.items():
|
||
|
|
if k not in data:
|
||
|
|
data[k] = def_val
|
||
|
|
else:
|
||
|
|
data[k] = replace_if_invalid(value=data[k], replacement=def_val, pattern=RE_VALID_OUTDIR)
|
||
|
|
|
||
|
|
# Check and fix other default settings
|
||
|
|
for k, def_val in DEFAULT_OTHER.items():
|
||
|
|
if k not in data:
|
||
|
|
data[k] = def_val
|
||
|
|
|
||
|
|
# Write results to file
|
||
|
|
dict_to_json_file(target_file or config_file, data)
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
if len(sys.argv) > 1:
|
||
|
|
check_and_replace_config(*sys.argv[1:])
|
||
|
|
else:
|
||
|
|
check_and_replace_config(DEFAULT_FILEPATH)
|
||
|
|
|