From 14c4b36affa4535dd58012574e39e227c2b3d603 Mon Sep 17 00:00:00 2001
From: AbdBarho <ka70911@gmail.com>
Date: Sat, 1 Oct 2022 12:57:53 +0200
Subject: [PATCH] v2 (#108)

### Update versions
- auto:
https://github.com/AUTOMATIC1111/stable-diffusion-webui/commit/3f417566b0bda8eab05d247567aebf001c1d1725

### Breaking changes:
* renamed `automatic-1111` service to `auto`
* the `cache` folder is now deprecated, replaced with `data` (see
migration guide below)
* `embeddings` folder has been moved to `data/embeddings`
* use GFPGAN 1.4

### Migration Guide

Note: in theory, running the command
```
docker compose --profile download up --build
```
is all you need to use the new version, however, this means you will
also have to download everything again. A new script is available under
`scripts/migratev1tov2.sh` that will copy models to the new structure
and should get you most of the way, run
```bash
./scripts/migratev1tov2.sh
```
or you can manually inspect the script and copy the files

After that, run
```
docker compose --profile download up --build
```
to validate everything.
---
 {scripts => .devscripts}/chmod.sh       |  0
 cache/.gitignore                        |  5 ---
 data/.gitignore                         | 13 ++++++
 docker-compose.yml                      |  7 ++--
 scripts/migratev1tov2.sh                | 30 ++++++++++++++
 services/AUTOMATIC1111/Dockerfile       | 10 ++---
 services/AUTOMATIC1111/mount.sh         | 53 +++++++++++--------------
 services/download/checksums.sha256      | 12 +++---
 services/download/download.sh           |  7 +++-
 services/download/links.txt             | 14 +++----
 services/hlky/Dockerfile                |  6 +--
 services/hlky/mount.sh                  | 39 ++++++++----------
 services/hlky/userconfig_streamlit.yaml |  2 +-
 13 files changed, 113 insertions(+), 85 deletions(-)
 rename {scripts => .devscripts}/chmod.sh (100%)
 delete mode 100644 cache/.gitignore
 create mode 100644 data/.gitignore
 create mode 100755 scripts/migratev1tov2.sh

diff --git a/scripts/chmod.sh b/.devscripts/chmod.sh
similarity index 100%
rename from scripts/chmod.sh
rename to .devscripts/chmod.sh
diff --git a/cache/.gitignore b/cache/.gitignore
deleted file mode 100644
index a2aafd4..0000000
--- a/cache/.gitignore
+++ /dev/null
@@ -1,5 +0,0 @@
-/torch
-/transformers
-/weights
-/models
-/custom-models
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..2674e83 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:
@@ -30,7 +30,7 @@ services:
     environment:
       - CLI_ARGS=--optimized-turbo
 
-  automatic1111: &automatic
+  auto: &automatic
     <<: *base_service
     profiles: ["auto"]
     build: ./services/AUTOMATIC1111
@@ -38,11 +38,10 @@ services:
       - *v1
       - *v2
       - ./services/AUTOMATIC1111/config.json:/stable-diffusion-webui/config.json
-      - ./embeddings:/stable-diffusion-webui/embeddings
     environment:
       - CLI_ARGS=--allow-code --medvram
 
-  automatic1111-cpu:
+  auto-cpu:
     <<: *automatic
     profiles: ["auto-cpu"]
     deploy: {}
diff --git a/scripts/migratev1tov2.sh b/scripts/migratev1tov2.sh
new file mode 100755
index 0000000..0b67158
--- /dev/null
+++ b/scripts/migratev1tov2.sh
@@ -0,0 +1,30 @@
+mkdir -p data/.cache data/StableDiffusion data/Codeformer data/GFPGAN data/ESRGAN data/BSRGAN data/RealESRGAN data/SwinIR data/LDSR /data/embeddings
+
+cp -vf cache/models/model.ckpt data/StableDiffusion/model.ckpt
+
+cp -vf cache/models/LDSR.ckpt data/LDSR/model.ckpt
+cp -vf cache/models/LDSR.yaml data/LDSR/project.yaml
+
+cp -vf cache/models/RealESRGAN_x4plus.pth data/RealESRGAN/
+cp -vf cache/models/RealESRGAN_x4plus_anime_6B.pth data/RealESRGAN/
+
+cp -vrf cache/torch data/.cache/
+
+mkdir -p data/.cache/huggingface/transformers/
+cp -vrf cache/transformers/* data/.cache/huggingface/transformers/
+
+cp -v cache/custom-models/* data/StableDiffusion/
+
+mkdir -p data/.cache/clip/
+cp -vf cache/weights/ViT-L-14.pt data/.cache/clip/
+
+cp -vf cache/weights/codeformer.pth data/Codeformer/codeformer-v0.1.0.pth
+
+cp -vf cache/weights/detection_Resnet50_Final.pth data/.cache/
+cp -vf cache/weights/parsing_parsenet.pth data/.cache/
+
+cp -v embeddings/* data/embeddings/
+
+echo this script was created 10/2022
+echo Dont forget to run: docker compose --profile download up --build
+echo the cache and embeddings folders can be deleted, but its not necessary.
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..2108f50 100755
--- a/services/AUTOMATIC1111/mount.sh
+++ b/services/AUTOMATIC1111/mount.sh
@@ -1,36 +1,31 @@
 #!/bin/bash
 
-set -e
+set -Eeuo pipefail
 
-declare -A MODELS
+declare -A MOUNTS
 
-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
+MOUNTS["/root/.cache"]="/data/.cache"
 
-MODELS_DIR=/cache/models
+# main
+MOUNTS["${ROOT}/models/Stable-diffusion"]="/data/StableDiffusion"
+MOUNTS["${ROOT}/models/Codeformer"]="/data/Codeformer"
+MOUNTS["${ROOT}/models/GFPGAN"]="/data/GFPGAN"
+MOUNTS["${ROOT}/models/ESRGAN"]="/data/ESRGAN"
+MOUNTS["${ROOT}/models/BSRGAN"]="/data/BSRGAN"
+MOUNTS["${ROOT}/models/RealESRGAN"]="/data/RealESRGAN"
+MOUNTS["${ROOT}/models/SwinIR"]="/data/SwinIR"
+MOUNTS["${ROOT}/models/LDSR"]="/data/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
+MOUNTS["${ROOT}/embeddings"]="/data/embeddings"
+
+# extra hacks
+MOUNTS["${ROOT}/repositories/CodeFormer/weights/facelib"]="/data/.cache"
+
+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 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
-
-mkdir -p /cache/torch /cache/transformers /cache/weights /cache/models /cache/custom-models
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..f979fa0 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,8 +13,11 @@ 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..."
 
 parallel --will-cite -a /docker/checksums.sha256 "echo -n {} | sha256sum -c"
+
+# fix potential permissions
+chmod -R 777 /data /output
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
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