Falln87 commited on
Commit
47e038b
·
verified ·
1 Parent(s): e1558bd

Create config.py

Browse files
Files changed (1) hide show
  1. config.py +77 -0
config.py ADDED
@@ -0,0 +1,77 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+
3
+ """Checks and sets default values for config.json before starting the container."""
4
+
5
+ import json
6
+ import re
7
+ import os.path
8
+ import sys
9
+
10
+ DEFAULT_FILEPATH = '/data/config/auto/config.json'
11
+
12
+ DEFAULT_OUTDIRS = {
13
+ "outdir_samples": "",
14
+ "outdir_txt2img_samples": "/output/txt2img",
15
+ "outdir_img2img_samples": "/output/img2img",
16
+ "outdir_extras_samples": "/output/extras",
17
+ "outdir_grids": "",
18
+ "outdir_txt2img_grids": "/output/txt2img-grids",
19
+ "outdir_img2img_grids": "/output/img2img-grids",
20
+ "outdir_save": "/output/saved",
21
+ "outdir_init_images": "/output/init-images",
22
+ }
23
+ RE_VALID_OUTDIR = re.compile(r"(^/output(/\.?[\w\-\_]+)+/?$)|(^\s?$)")
24
+
25
+ DEFAULT_OTHER = {
26
+ "font": "DejaVuSans.ttf",
27
+ }
28
+
29
+ def dict_to_json_file(target_file: str, data: dict):
30
+ """Write dictionary to specified json file"""
31
+
32
+ with open(target_file, 'w') as f:
33
+ json.dump(data, f)
34
+
35
+ def json_file_to_dict(config_file: str) -> dict|None:
36
+ """Load json file into a dictionary. Return None if file does not exist."""
37
+
38
+ if os.path.isfile(config_file):
39
+ with open(config_file, 'r') as f:
40
+ return json.load(f)
41
+ else:
42
+ return None
43
+
44
+ def replace_if_invalid(value: str, replacement: str, pattern: str|re.Pattern[str]) -> str:
45
+ """Returns original value if valid, fallback value if invalid"""
46
+
47
+ if re.match(pattern, value):
48
+ return value
49
+ else:
50
+ return replacement
51
+
52
+ def check_and_replace_config(config_file: str, target_file: str = None):
53
+ """Checks given file for invalid values. Replaces those with fallback values (default: overwrites file)."""
54
+
55
+ # Get current user config, or empty if file does not exists
56
+ data = json_file_to_dict(config_file) or {}
57
+
58
+ # Check and fix output directories
59
+ for k, def_val in DEFAULT_OUTDIRS.items():
60
+ if k not in data:
61
+ data[k] = def_val
62
+ else:
63
+ data[k] = replace_if_invalid(value=data[k], replacement=def_val, pattern=RE_VALID_OUTDIR)
64
+
65
+ # Check and fix other default settings
66
+ for k, def_val in DEFAULT_OTHER.items():
67
+ if k not in data:
68
+ data[k] = def_val
69
+
70
+ # Write results to file
71
+ dict_to_json_file(target_file or config_file, data)
72
+
73
+ if __name__ == '__main__':
74
+ if len(sys.argv) > 1:
75
+ check_and_replace_config(*sys.argv[1:])
76
+ else:
77
+ check_and_replace_config(DEFAULT_FILEPATH)