From ffdb9d23671fda6f00db58ca2b0d7052b63a6678 Mon Sep 17 00:00:00 2001
From: Abdullah Barhoum <ka70911@gmail.com>
Date: Fri, 16 Sep 2022 16:24:49 +0200
Subject: [PATCH] Update lstein to dev branch

---
 docker-compose.yml         |  3 +++
 services/lstein/Dockerfile | 25 +++++++++++++++++++------
 services/lstein/info.py    | 10 ++++++++++
 services/lstein/mount.sh   | 26 ++++++++++++++++++++++++++
 4 files changed, 58 insertions(+), 6 deletions(-)
 create mode 100644 services/lstein/info.py
 create mode 100644 services/lstein/mount.sh

diff --git a/docker-compose.yml b/docker-compose.yml
index 5d284e3..4742984 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -52,3 +52,6 @@ services:
     <<: *base_service
     profiles: ["lstein"]
     build: ./services/lstein/
+    environment:
+      - PRELOAD=false
+      - CLI_ARGS=
diff --git a/services/lstein/Dockerfile b/services/lstein/Dockerfile
index f3b7c6f..f998bc6 100644
--- a/services/lstein/Dockerfile
+++ b/services/lstein/Dockerfile
@@ -20,12 +20,25 @@ conda env update --file environment.yaml -n base
 conda clean -a -y
 EOF
 
-ENV TRANSFORMERS_CACHE=/cache/transformers TORCH_HOME=/cache/torch CLI_ARGS=""
 
+ARG BRANCH=development SHA=45af30f3a4c98b50c755717831c5fff75a3a8b43
+# ARG BRANCH=main SHA=89da371f4841f7e05da5a1672459d700c3920784
+RUN <<EOF
+cd stable-diffusion
+git fetch
+git checkout ${BRANCH}
+git reset --hard ${SHA}
+conda env update --file environment.yaml -n base
+conda clean -a -y
+EOF
+
+RUN pip uninstall opencv-python -y && pip install --prefer-binary --upgrade --force-reinstall --no-cache-dir opencv-python-headless
+
+COPY . /docker/
+RUN python3 /docker/info.py /stable-diffusion/static/dream_web/index.html
+
+ENV TRANSFORMERS_CACHE=/cache/transformers TORCH_HOME=/cache/torch PRELOAD=false CLI_ARGS=""
 WORKDIR /stable-diffusion
-
 EXPOSE 7860
-# run, -u to not buffer stdout / stderr
-CMD mkdir -p /stable-diffusion/models/ldm/stable-diffusion-v1/ && \
-  ln -sf /cache/models/model.ckpt /stable-diffusion/models/ldm/stable-diffusion-v1/model.ckpt && \
-  python3 -u scripts/dream.py --outdir /output --web --host 0.0.0.0 --port 7860 ${CLI_ARGS}
+
+CMD /docker/mount.sh && python3 -u scripts/dream.py --outdir /output --web --host 0.0.0.0 --port 7860 ${CLI_ARGS}
diff --git a/services/lstein/info.py b/services/lstein/info.py
new file mode 100644
index 0000000..1d8f476
--- /dev/null
+++ b/services/lstein/info.py
@@ -0,0 +1,10 @@
+import sys
+from pathlib import Path
+
+file = Path(sys.argv[1])
+file.write_text(
+  file.read_text()\
+  .replace('GitHub site</a>', """
+  GitHub site</a>, Deployed with <a href="https://github.com/AbdBarho/stable-diffusion-webui-docker/">stable-diffusion-webui-docker</a>
+""", 1)
+)
diff --git a/services/lstein/mount.sh b/services/lstein/mount.sh
new file mode 100644
index 0000000..ef36444
--- /dev/null
+++ b/services/lstein/mount.sh
@@ -0,0 +1,26 @@
+#!/bin/bash
+
+set -eu
+
+ROOT=/stable-diffusion
+
+mkdir -p "${ROOT}/models/ldm/stable-diffusion-v1/"
+ln -sf /cache/models/model.ckpt "${ROOT}/models/ldm/stable-diffusion-v1/model.ckpt"
+
+if test -f /cache/models/GFPGANv1.3.pth; then
+  base="${ROOT}/src/gfpgan/experiments/pretrained_models/"
+  mkdir -p "${base}"
+  ln -sf /cache/models/GFPGANv1.3.pth "${base}/GFPGANv1.3.pth"
+  echo "Mounted GFPGANv1.3.pth"
+fi
+
+# facexlib
+FACEX_WEIGHTS=/opt/conda/lib/python3.8/site-packages/facexlib/weights
+
+rm -rf "${FACEX_WEIGHTS}"
+mkdir -p /cache/weights
+ln -sf -T /cache/weights "${FACEX_WEIGHTS}"
+
+if "${PRELOAD}" == "true"; then
+  python3 -u scripts/preload_models.py
+fi