From bb6b6386f03236a761275bd285c1fbdde9052d2c Mon Sep 17 00:00:00 2001
From: AbdBarho <ka70911@gmail.com>
Date: Sat, 1 Oct 2022 12:17:53 +0200
Subject: [PATCH] migrate hlky

---
 services/download/download.sh           |  3 ++
 services/hlky/Dockerfile                |  6 ++--
 services/hlky/mount.sh                  | 39 +++++++++++--------------
 services/hlky/userconfig_streamlit.yaml |  2 +-
 4 files changed, 24 insertions(+), 26 deletions(-)

diff --git a/services/download/download.sh b/services/download/download.sh
index 16f0c92..f979fa0 100755
--- a/services/download/download.sh
+++ b/services/download/download.sh
@@ -18,3 +18,6 @@ aria2c --input-file /docker/links.txt --dir /data --continue
 echo "Checking SHAs..."
 
 parallel --will-cite -a /docker/checksums.sha256 "echo -n {} | sha256sum -c"
+
+# fix potential permissions
+chmod -R 777 /data /output
diff --git a/services/hlky/Dockerfile b/services/hlky/Dockerfile
index c0fa6b1..7fd12d5 100644
--- a/services/hlky/Dockerfile
+++ b/services/hlky/Dockerfile
@@ -42,9 +42,9 @@ COPY . /docker/
 RUN python /docker/info.py /stable-diffusion/frontend/frontend.py && chmod +x /docker/mount.sh
 
 WORKDIR /stable-diffusion
-ENV TRANSFORMERS_CACHE=/cache/transformers TORCH_HOME=/cache/torch PYTHONPATH="${PYTHONPATH}:${PWD}" CLI_ARGS=""
+ENV PYTHONPATH="${PYTHONPATH}:${PWD}" CLI_ARGS=""
 EXPOSE 7860
 # run, -u to not buffer stdout / stderr
 CMD /docker/mount.sh && \
-  python3 -u scripts/webui.py --outdir /output --ckpt /cache/models/model.ckpt ${CLI_ARGS}
-# sed -i -- 's/8501/7860/g' .streamlit/config.toml && STREAMLIT_SERVER_HEADLESS=true python -u -m streamlit run scripts/webui_streamlit.py  --theme.base dark --server.port 7860
+  python3 -u scripts/webui.py --outdir /output --ckpt /data/StableDiffusion/model.ckpt ${CLI_ARGS}
+# sed -i -- 's/8501/7860/g' .streamlit/config.toml && STREAMLIT_SERVER_HEADLESS=true python -u -m streamlit run scripts/webui_streamlit.py  --theme.base dark
diff --git a/services/hlky/mount.sh b/services/hlky/mount.sh
index 39d1d4f..95e7f13 100755
--- a/services/hlky/mount.sh
+++ b/services/hlky/mount.sh
@@ -1,33 +1,28 @@
 #!/bin/bash
 
-set -e
+set -Eeuo pipefail
 
-declare -A MODELS
+declare -A MOUNTS
 
 ROOT=/stable-diffusion/src
 
-MODELS["${ROOT}/gfpgan/experiments/pretrained_models/GFPGANv1.3.pth"]=GFPGANv1.3.pth
-MODELS["${ROOT}/realesrgan/experiments/pretrained_models/RealESRGAN_x4plus.pth"]=RealESRGAN_x4plus.pth
-MODELS["${ROOT}/realesrgan/experiments/pretrained_models/RealESRGAN_x4plus_anime_6B.pth"]=RealESRGAN_x4plus_anime_6B.pth
-MODELS["${ROOT}/latent-diffusion/experiments/pretrained_models/model.ckpt"]=LDSR.ckpt
-MODELS["${ROOT}/latent-diffusion/experiments/pretrained_models/project.yaml"]=LDSR.yaml
+# cache
+MOUNTS["/root/.cache"]=/data/.cache
+# ui specific
+MOUNTS["${ROOT}/gfpgan/experiments/pretrained_models/GFPGANv1.3.pth"]=/data/GFPGAN/GFPGANv1.4.pth
+MOUNTS["${ROOT}/realesrgan/experiments/pretrained_models"]=/data/RealESRGAN
+MOUNTS["${ROOT}/latent-diffusion/experiments/pretrained_models"]=/data/LDSR
+# hacks
+MOUNTS["/stable-diffusion/gfpgan/weights"]=/data/.cache
 
-MODELS_DIR=/cache/models
-
-for path in "${!MODELS[@]}"; do
-  name=${MODELS[$path]}
-  base=$(dirname "${path}")
-  from_path="${MODELS_DIR}/${name}"
-  if test -f "${from_path}"; then
-    mkdir -p "${base}" && ln -sf "${from_path}" "${path}" && echo "Mounted ${name}"
-  else
-    echo "Skipping ${name}"
-  fi
+for to_path in "${!MOUNTS[@]}"; do
+  set -Eeuo pipefail
+  from_path="${MOUNTS[${to_path}]}"
+  rm -rf "${to_path}"
+  mkdir -p "$(dirname "${to_path}")"
+  ln -sT "${from_path}" "${to_path}"
+  echo Mounted $(basename "${from_path}")
 done
 
-# force facexlib cache
-mkdir -p /cache/weights/ /stable-diffusion/gfpgan/
-ln -sf /cache/weights/ /stable-diffusion/gfpgan/
-
 # streamlit config
 ln -sf /docker/userconfig_streamlit.yaml /stable-diffusion/configs/webui/userconfig_streamlit.yaml
diff --git a/services/hlky/userconfig_streamlit.yaml b/services/hlky/userconfig_streamlit.yaml
index 51fa016..80b59fb 100644
--- a/services/hlky/userconfig_streamlit.yaml
+++ b/services/hlky/userconfig_streamlit.yaml
@@ -1,7 +1,7 @@
 general:
   outdir: /outputs
   default_model: "Stable Diffusion v1.4"
-  default_model_path: /cache/models/model.ckpt
+  default_model_path: /data/StableDiffusion/model.ckpt
   outdir_txt2img: /outputs/txt2img-samples
   outdir_img2img: /outputs/img2img-samples
   optimized: True