From aae83bb8f2d24f1dcd9b7d77fbd32b9b69f0927e Mon Sep 17 00:00:00 2001
From: AbdBarho <ka70911@gmail.com>
Date: Fri, 16 Sep 2022 16:40:20 +0200
Subject: [PATCH] Update lstein to dev branch (#73)

---
 README.md                  |  7 ++++++-
 docker-compose.yml         |  3 +++
 services/lstein/Dockerfile | 25 +++++++++++++++++++------
 services/lstein/info.py    | 10 ++++++++++
 services/lstein/mount.sh   | 26 ++++++++++++++++++++++++++
 5 files changed, 64 insertions(+), 7 deletions(-)
 create mode 100644 services/lstein/info.py
 create mode 100644 services/lstein/mount.sh

diff --git a/README.md b/README.md
index 4abe75f..5a2faca 100644
--- a/README.md
+++ b/README.md
@@ -41,13 +41,18 @@ Screenshots:
 
 ### lstein
 
-[lstein's fork](https://github.com/lstein/stable-diffusion) is very mature when it comes to the cli, but less so for the WebUI.
+[lstein's fork](https://github.com/lstein/stable-diffusion) is very mature when it comes to the cli, and the WebUI has potential.
+
+| Text to image                                                                                              | Image to image                                                                                             | Extras                                                                                                     |
+| ---------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------- |
+| ![](https://user-images.githubusercontent.com/24505302/190662506-dabdc967-93af-4d78-8533-394604d29ba4.jpg) | ![](https://user-images.githubusercontent.com/24505302/190662557-7640d9f0-30d8-4527-97b0-07d3f48108d4.jpg) | ![](https://user-images.githubusercontent.com/24505302/190662588-37a01fad-f993-4674-9ae6-8714aa229f7b.jpg) |
 
 ## Setup & Usage
 
 Visit the wiki for [Setup](https://github.com/AbdBarho/stable-diffusion-webui-docker/wiki/Setup) and [Usage](https://github.com/AbdBarho/stable-diffusion-webui-docker/wiki/Usage) instructions, checkout the [FAQ](https://github.com/AbdBarho/stable-diffusion-webui-docker/wiki/FAQ) page if you face any problems, or create a new issue!
 
 ## Contributing
+
 Contributions are welcome! create an issue first of what you want to contribute (before you implement anything) so we can talk about it.
 
 ## Disclaimer
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