From 5b547a515db7f632aa0d283d6c160327804cdc9a Mon Sep 17 00:00:00 2001
From: AbdBarho <ka70911@gmail.com>
Date: Sat, 1 Oct 2022 10:53:48 +0200
Subject: [PATCH] Restructure data + auto

---
 data/.gitignore                    | 13 +++++
 docker-compose.yml                 |  2 +-
 services/AUTOMATIC1111/Dockerfile  | 10 ++--
 services/AUTOMATIC1111/mount.sh    | 77 ++++++++++++++++++++----------
 services/download/checksums.sha256 | 12 ++---
 services/download/download.sh      |  4 +-
 services/download/links.txt        | 14 +++---
 7 files changed, 84 insertions(+), 48 deletions(-)
 create mode 100644 data/.gitignore

diff --git a/data/.gitignore b/data/.gitignore
new file mode 100644
index 0000000..d8caa8b
--- /dev/null
+++ b/data/.gitignore
@@ -0,0 +1,13 @@
+# for all of the stuff downloaded by transformers, pytorch, and others
+/.cache
+# for all stable diffusion models (main, waifu diffusion, etc..)
+/StableDiffusion
+# others
+/Codeformer
+/GFPGAN
+/ESRGAN
+/BSRGAN
+/RealESRGAN
+/SwinIR
+/LDSR
+# TODO: add the embeddings folder here
diff --git a/docker-compose.yml b/docker-compose.yml
index e3f80b1..a5d8419 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -4,7 +4,7 @@ x-base_service: &base_service
     ports:
       - "7860:7860"
     volumes:
-      - &v1 ./cache:/cache
+      - &v1 ./data:/data
       - &v2 ./output:/output
     deploy:
       resources:
diff --git a/services/AUTOMATIC1111/Dockerfile b/services/AUTOMATIC1111/Dockerfile
index e11d0a5..f6a70a7 100644
--- a/services/AUTOMATIC1111/Dockerfile
+++ b/services/AUTOMATIC1111/Dockerfile
@@ -49,7 +49,7 @@ RUN pip install --prefer-binary --no-cache-dir -r ${ROOT}/repositories/CodeForme
 # Note: don't update the sha of previous versions because the install will take forever
 # instead, update the repo state in a later step
 
-ARG SHA=f80c3696f63a181f720105559d42ee53453ed0eb
+ARG SHA=3f417566b0bda8eab05d247567aebf001c1d1725
 RUN <<EOF
 cd stable-diffusion-webui
 git pull --rebase
@@ -60,18 +60,16 @@ EOF
 
 RUN pip install --prefer-binary -U --no-cache-dir opencv-python-headless
 
-ENV TRANSFORMERS_CACHE=/cache/transformers TORCH_HOME=/cache/torch CLI_ARGS=""
 
 COPY . /docker
 RUN <<EOF
 chmod +x /docker/mount.sh && python3 /docker/info.py ${ROOT}/modules/ui.py
-# hackiest of hacks, change default cache dir of clip #88
-# https://github.com/openai/CLIP/blob/d50d76daa670286dd6cacf3bcd80b5e4823fc8e1/clip/clip.py#L94
-sed -i -- 's/download_root: str = None/download_root: str = "\/cache\/weights"/' /opt/conda/lib/python3.8/site-packages/clip/clip.py
 EOF
 
+ENV CLI_ARGS=""
 WORKDIR ${WORKDIR}
 EXPOSE 7860
 # run, -u to not buffer stdout / stderr
 CMD /docker/mount.sh && \
-  python3 -u ../../webui.py --listen --port 7860 --hide-ui-dir-config --ckpt-dir /cache/custom-models --ckpt /cache/models/model.ckpt --gfpgan-model /cache/models/GFPGANv1.3.pth ${CLI_ARGS}
+  python3 -u ../../webui.py --listen --port 7860 --hide-ui-dir-config --ckpt-dir ${ROOT}/models/Stable-diffusion ${CLI_ARGS}
+#  --gfpgan-models-path /data/GFPGAN --codeformer-models-path /data/Codeformer --esrgan-models-path /data/ESRGAN --bsrgan-models-path /data/BSRGAN --realesrgan-models-path /data/RealESRGAN --swinir-models-path /data/SwinIR --ldsr-models-path /data/LDSR
diff --git a/services/AUTOMATIC1111/mount.sh b/services/AUTOMATIC1111/mount.sh
index 671edc6..13f4f4a 100755
--- a/services/AUTOMATIC1111/mount.sh
+++ b/services/AUTOMATIC1111/mount.sh
@@ -1,36 +1,61 @@
 #!/bin/bash
 
-set -e
+set -Eeuo pipefail
 
-declare -A MODELS
+rm -rf /root/.cache
+ln -sT /data/.cache /root/.cache
 
-MODELS["${ROOT}/GFPGANv1.3.pth"]=GFPGANv1.3.pth
-MODELS["${WORKDIR}/repositories/latent-diffusion/experiments/pretrained_models/model.chkpt"]=LDSR.ckpt
-MODELS["${WORKDIR}/repositories/latent-diffusion/experiments/pretrained_models/project.yaml"]=LDSR.yaml
+declare -A FOLDERS
 
-MODELS_DIR=/cache/models
+# main
+FOLDERS["/data/StableDiffusion"]="${ROOT}/models/Stable-diffusion"
+FOLDERS["/data/Codeformer"]="${ROOT}/models/Codeformer"
+FOLDERS["/data/GFPGAN"]="${ROOT}/models/GFPGAN"
+FOLDERS["/data/ESRGAN"]="${ROOT}/models/ESRGAN"
+FOLDERS["/data/BSRGAN"]="${ROOT}/models/BSRGAN"
+FOLDERS["/data/RealESRGAN"]="${ROOT}/models/RealESRGAN"
+FOLDERS["/data/SwinIR"]="${ROOT}/models/SwinIR"
+FOLDERS["/data/LDSR"]="${ROOT}/models/LDSR"
 
-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
+# extra hacks
+FOLDERS["/data/.cache"]="${ROOT}/repositories/CodeFormer/weights/facelib"
+
+for from_path in "${!FOLDERS[@]}"; do
+  set -Eeuo pipefail
+  to_path="${FOLDERS[${from_path}]}"
+  rm -rf "${to_path}"
+  mkdir -p "$(dirname "${to_path}")"
+  ln -sT "${from_path}" "${to_path}"
+  echo Mounted $(basename "${from_path}")
 done
 
-# force realesrgan cache
-rm -rf /opt/conda/lib/python3.8/site-packages/realesrgan/weights
-ln -s -T "${MODELS_DIR}" /opt/conda/lib/python3.8/site-packages/realesrgan/weights
+# declare -A MODELS
 
-# force facexlib cache
-mkdir -p /cache/weights/ ${WORKDIR}/gfpgan/
-ln -sf /cache/weights/ ${WORKDIR}/gfpgan/
-# code former cache
-rm -rf ${ROOT}/repositories/CodeFormer/weights/CodeFormer ${ROOT}/repositories/CodeFormer/weights/facelib
-ln -sf -T /cache/weights ${ROOT}/repositories/CodeFormer/weights/CodeFormer
-ln -sf -T /cache/weights ${ROOT}/repositories/CodeFormer/weights/facelib
+# MODELS["${ROOT}/GFPGANv1.3.pth"]=GFPGANv1.3.pth
+# MODELS["${WORKDIR}/repositories/latent-diffusion/experiments/pretrained_models/model.chkpt"]=LDSR.ckpt
+# MODELS["${WORKDIR}/repositories/latent-diffusion/experiments/pretrained_models/project.yaml"]=LDSR.yaml
 
-mkdir -p /cache/torch /cache/transformers /cache/weights /cache/models /cache/custom-models
+# 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
+# done
+
+# # force realesrgan cache
+# rm -rf /opt/conda/lib/python3.8/site-packages/realesrgan/weights
+# ln -s -T "${MODELS_DIR}" /opt/conda/lib/python3.8/site-packages/realesrgan/weights
+
+# # force facexlib cache
+# mkdir -p /cache/weights/ ${WORKDIR}/gfpgan/
+# ln -sf /cache/weights/ ${WORKDIR}/gfpgan/
+# # code former cache
+# rm -rf ${ROOT}/repositories/CodeFormer/weights/CodeFormer ${ROOT}/repositories/CodeFormer/weights/facelib
+# ln -sf -T /cache/weights ${ROOT}/repositories/CodeFormer/weights/CodeFormer
+# ln -sf -T /cache/weights ${ROOT}/repositories/CodeFormer/weights/facelib
diff --git a/services/download/checksums.sha256 b/services/download/checksums.sha256
index 4cf6de3..a83a4a0 100644
--- a/services/download/checksums.sha256
+++ b/services/download/checksums.sha256
@@ -1,6 +1,6 @@
-fe4efff1e174c627256e44ec2991ba279b3816e364b49f9be2abc0b3ff3f8556  /cache/models/model.ckpt
-c953a88f2727c85c3d9ae72e2bd4846bbaf59fe6972ad94130e23e7017524a70  /cache/models/GFPGANv1.3.pth
-4fa0d38905f75ac06eb49a7951b426670021be3018265fd191d2125df9d682f1  /cache/models/RealESRGAN_x4plus.pth
-f872d837d3c90ed2e05227bed711af5671a6fd1c9f7d7e91c911a61f155e99da  /cache/models/RealESRGAN_x4plus_anime_6B.pth
-c209caecac2f97b4bb8f4d726b70ac2ac9b35904b7fc99801e1f5e61f9210c13  /cache/models/LDSR.ckpt
-9d6ad53c5dafeb07200fb712db14b813b527edd262bc80ea136777bdb41be2ba  /cache/models/LDSR.yaml
+fe4efff1e174c627256e44ec2991ba279b3816e364b49f9be2abc0b3ff3f8556  /data/StableDiffusion/model.ckpt
+e2cd4703ab14f4d01fd1383a8a8b266f9a5833dacee8e6a79d3bf21a1b6be5ad  /data/GFPGAN/GFPGANv1.4.pth
+4fa0d38905f75ac06eb49a7951b426670021be3018265fd191d2125df9d682f1  /data/RealESRGAN/RealESRGAN_x4plus.pth
+f872d837d3c90ed2e05227bed711af5671a6fd1c9f7d7e91c911a61f155e99da  /data/RealESRGAN/RealESRGAN_x4plus_anime_6B.pth
+c209caecac2f97b4bb8f4d726b70ac2ac9b35904b7fc99801e1f5e61f9210c13  /data/LDSR/model.ckpt
+9d6ad53c5dafeb07200fb712db14b813b527edd262bc80ea136777bdb41be2ba  /data/LDSR/project.yaml
diff --git a/services/download/download.sh b/services/download/download.sh
index 394b805..16f0c92 100755
--- a/services/download/download.sh
+++ b/services/download/download.sh
@@ -2,7 +2,7 @@
 
 set -Eeuo pipefail
 
-mkdir -p /cache/torch /cache/transformers /cache/weights /cache/models /cache/custom-models
+mkdir -p /data/.cache /data/StableDiffusion /data/Codeformer /data/GFPGAN /data/ESRGAN /data/BSRGAN /data/RealESRGAN /data/SwinIR /data/LDSR
 
 cat <<EOF
 By using this software, you agree to the following licenses:
@@ -13,7 +13,7 @@ EOF
 
 echo "Downloading, this might take a while..."
 
-aria2c --input-file /docker/links.txt --dir /cache/models --continue
+aria2c --input-file /docker/links.txt --dir /data --continue
 
 echo "Checking SHAs..."
 
diff --git a/services/download/links.txt b/services/download/links.txt
index b728297..faafa33 100644
--- a/services/download/links.txt
+++ b/services/download/links.txt
@@ -1,12 +1,12 @@
 https://www.googleapis.com/storage/v1/b/aai-blog-files/o/sd-v1-4.ckpt?alt=media
-  out=model.ckpt
-https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth
-  out=GFPGANv1.3.pth
+  out=StableDiffusion/model.ckpt
+https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.4.pth
+  out=GFPGAN/GFPGANv1.4.pth
 https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth
-  out=RealESRGAN_x4plus.pth
+  out=RealESRGAN/RealESRGAN_x4plus.pth
 https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth
-  out=RealESRGAN_x4plus_anime_6B.pth
+  out=RealESRGAN/RealESRGAN_x4plus_anime_6B.pth
 https://heibox.uni-heidelberg.de/f/31a76b13ea27482981b4/?dl=1
-  out=LDSR.yaml
+  out=LDSR/project.yaml
 https://heibox.uni-heidelberg.de/f/578df07c8fc04ffbadf3/?dl=1
-  out=LDSR.ckpt
+  out=LDSR/model.ckpt