From e599fce0edc5bd639688e88dddbe43bc7060ddfa Mon Sep 17 00:00:00 2001
From: Mat <m.falkowski@dotmat.de>
Date: Thu, 4 May 2023 22:57:05 +0200
Subject: [PATCH] jq query to python script, gitignore whitelist

---
 .gitignore                           |  4 ++
 data/.gitignore                      | 31 ++---------
 services/AUTOMATIC1111/config.json   | 11 ----
 services/AUTOMATIC1111/config.py     | 78 ++++++++++++++++++++++++++++
 services/AUTOMATIC1111/entrypoint.sh | 18 +------
 5 files changed, 89 insertions(+), 53 deletions(-)
 delete mode 100644 services/AUTOMATIC1111/config.json
 create mode 100644 services/AUTOMATIC1111/config.py

diff --git a/.gitignore b/.gitignore
index f64b472..340919c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,6 @@
 /.devcontainer
 /docker-compose.override.yml
+
+# VSCode specific
+*.code-workspace
+/.vscode/settings.json
diff --git a/data/.gitignore b/data/.gitignore
index 9649426..97a7f76 100644
--- a/data/.gitignore
+++ b/data/.gitignore
@@ -1,26 +1,5 @@
-# for all of the stuff downloaded by transformers, pytorch, and others
-/.cache
-# for UIs
-/config
-# for all stable diffusion models (main, waifu diffusion, etc..)
-/StableDiffusion
-# others
-/Codeformer
-/GFPGAN
-/ESRGAN
-/BSRGAN
-/RealESRGAN
-/SwinIR
-/MiDaS
-/BLIP
-/ScuNET
-/LDSR
-/Deepdanbooru
-/Hypernetworks
-/VAE
-/embeddings
-/Lora
-/ControlNet
-/openpose
-/ModelScope
-/LyCORIS
+# Ignore all
+*
+
+# Whitelist
+!/.gitignore
diff --git a/services/AUTOMATIC1111/config.json b/services/AUTOMATIC1111/config.json
deleted file mode 100644
index e02aa35..0000000
--- a/services/AUTOMATIC1111/config.json
+++ /dev/null
@@ -1,11 +0,0 @@
-{
-  "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",
-  "font": "DejaVuSans.ttf"
-}
diff --git a/services/AUTOMATIC1111/config.py b/services/AUTOMATIC1111/config.py
new file mode 100644
index 0000000..b8d699e
--- /dev/null
+++ b/services/AUTOMATIC1111/config.py
@@ -0,0 +1,78 @@
+#!/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)
+
diff --git a/services/AUTOMATIC1111/entrypoint.sh b/services/AUTOMATIC1111/entrypoint.sh
index 696809c..54e709c 100755
--- a/services/AUTOMATIC1111/entrypoint.sh
+++ b/services/AUTOMATIC1111/entrypoint.sh
@@ -8,22 +8,8 @@ mkdir -p /data/config/auto/scripts/
 find "${ROOT}/scripts/" -maxdepth 1 -type l -delete
 cp -vrfTs /data/config/auto/scripts/ "${ROOT}/scripts/"
 
-# config.json for auto
-## Step 1: Copy default config to user config if missing
-cp -n /docker/config.json /data/config/auto/config.json
-## Step 2: Put the default config and user config into vars
-cfg_json_default=$(cat /docker/config.json)
-cfg_json_user=$(cat /data/config/auto/config.json)
-## Step 3: delete invalid entries from user.json
-cfg_json_user=$(echo $cfg_json_user | jq '
-  . as $original
-    | reduce (to_entries[] | select((.key | startswith("outdir_"))
-        and (.value | test("^\/output(\\.)?(\/\\.?[\\w\\-\\_]+)+\/?")|not)))
-        as $item
-        ($original; del(.[$item.key]))
-')
-## Step 4: Merge and save config
-echo $cfg_json_default $cfg_json_user | jq '. * input' | sponge /data/config/auto/config.json
+# Set up config file
+python /docker/config.py /data/config/auto/config.json
 
 if [ ! -f /data/config/auto/ui-config.json ]; then
   echo '{}' >/data/config/auto/ui-config.json