From 1cfb915d1259bd50d064dd6fa5366c8a19d18085 Mon Sep 17 00:00:00 2001
From: AbdBarho <ka70911@gmail.com>
Date: Fri, 4 Nov 2022 22:35:44 +0100
Subject: [PATCH] Invoke AI v2.1 (#195)

https://github.com/invoke-ai/InvokeAI/commit/6b89adfa7ebd4591ed91a76ecf152e3517cda385
---
 data/.gitignore                             |  1 +
 docker-compose.yml                          |  4 ++--
 services/download/checksums.sha256          |  1 +
 services/download/links.txt                 |  2 ++
 services/lstein/Dockerfile                  | 17 +++++++++--------
 services/lstein/{mount.sh => entrypoint.sh} |  6 ++++--
 services/lstein/models.yaml                 | 15 +++++++++++++++
 7 files changed, 34 insertions(+), 12 deletions(-)
 rename services/lstein/{mount.sh => entrypoint.sh} (85%)
 mode change 100755 => 100644
 create mode 100644 services/lstein/models.yaml

diff --git a/data/.gitignore b/data/.gitignore
index 36b7867..d4efe83 100644
--- a/data/.gitignore
+++ b/data/.gitignore
@@ -14,4 +14,5 @@
 /ScuNET
 /LDSR
 /Hypernetworks
+/VAE
 /embeddings
diff --git a/docker-compose.yml b/docker-compose.yml
index 1422685..4f54d3f 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -51,7 +51,7 @@ services:
     <<: *base_service
     profiles: ["lstein"]
     build: ./services/lstein/
-    image: sd-lstein:4
+    image: sd-lstein:5
     environment:
       - PRELOAD=true
-      - CLI_ARGS=
+      - CLI_ARGS=--max_loaded_models=1
diff --git a/services/download/checksums.sha256 b/services/download/checksums.sha256
index a83a4a0..5473ac9 100644
--- a/services/download/checksums.sha256
+++ b/services/download/checksums.sha256
@@ -4,3 +4,4 @@ e2cd4703ab14f4d01fd1383a8a8b266f9a5833dacee8e6a79d3bf21a1b6be5ad  /data/GFPGAN/G
 f872d837d3c90ed2e05227bed711af5671a6fd1c9f7d7e91c911a61f155e99da  /data/RealESRGAN/RealESRGAN_x4plus_anime_6B.pth
 c209caecac2f97b4bb8f4d726b70ac2ac9b35904b7fc99801e1f5e61f9210c13  /data/LDSR/model.ckpt
 9d6ad53c5dafeb07200fb712db14b813b527edd262bc80ea136777bdb41be2ba  /data/LDSR/project.yaml
+c6a580b13a5bc05a5e16e4dbb80608ff2ec251a162311590c1f34c013d7f3dab  /data/VAE/vae-ft-mse-840000-ema-pruned.ckpt
diff --git a/services/download/links.txt b/services/download/links.txt
index fbf20bf..4c03868 100644
--- a/services/download/links.txt
+++ b/services/download/links.txt
@@ -8,6 +8,8 @@ magnet:?xt=urn:btih:2daef5b5f63a16a9af9169a529b1a773fc452637&dn=v1-5-pruned-emao
   check-integrity=true
 https://drive.yerf.org/wl/?id=EBfTrmcCCUAGaQBXVIj5lJmEhjoP1tgl&mode=grid&download=1
   out=StableDiffusion/model.ckpt
+https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt
+  out=VAE/vae-ft-mse-840000-ema-pruned.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
diff --git a/services/lstein/Dockerfile b/services/lstein/Dockerfile
index 6010a77..3746a8b 100644
--- a/services/lstein/Dockerfile
+++ b/services/lstein/Dockerfile
@@ -8,21 +8,22 @@ ENV DEBIAN_FRONTEND=noninteractive PIP_EXISTS_ACTION=w  PIP_PREFER_BINARY=1 PIP_
 
 # now it requires python3.9
 RUN conda install python=3.9 && conda clean -a -y
-RUN conda install pytorch==1.11.0 torchvision==0.12.0 cudatoolkit=11.3 -c pytorch && conda clean -a -y
+RUN conda install cudatoolkit=11.6 -c conda-forge && conda clean -a -y
+RUN conda install pytorch==1.12.1 -c pytorch && conda clean -a -y
 
 RUN apt-get update && apt install fonts-dejavu-core -y && apt-get clean
 
 RUN <<EOF
 git clone https://github.com/invoke-ai/InvokeAI.git stable-diffusion
 cd stable-diffusion
-git reset --hard 986cc49d7a8441a70c709c04a55016a1c783f2ed
+git reset --hard 6b89adfa7ebd4591ed91a76ecf152e3517cda385
 git config --global http.postBuffer 1048576000
-# conda env update --file environment.yml -n base
+conda env update --file environment.yml -n base
 conda clean -a -y
 EOF
 
 
-ARG BRANCH=development SHA=fdf9b1c40c243d6806e787877e42ae4c115790ee
+ARG BRANCH=main SHA=6b89adfa7ebd4591ed91a76ecf152e3517cda385
 RUN <<EOF
 cd stable-diffusion
 git fetch
@@ -33,12 +34,11 @@ conda env update --file environment.yml -n base
 conda clean -a -y
 EOF
 
-RUN pip uninstall opencv-python -y && pip install --prefer-binary --force-reinstall --no-cache-dir opencv-python-headless
+RUN pip uninstall opencv-python -y && pip install --force-reinstall opencv-python-headless==4.5.5.64
 
 COPY . /docker/
 RUN <<EOF
 python3 /docker/info.py /stable-diffusion/frontend/dist/index.html
-chmod +x /docker/mount.sh
 EOF
 
 
@@ -46,5 +46,6 @@ ENV PRELOAD=false CLI_ARGS=""
 WORKDIR /stable-diffusion
 EXPOSE 7860
 
-CMD /docker/mount.sh && \
-  python3 -u scripts/invoke.py --outdir /output --web --host 0.0.0.0 --port 7860 ${CLI_ARGS}
+
+ENTRYPOINT ["/docker/entrypoint.sh"]
+CMD python3 -u scripts/invoke.py --outdir /output --web --host 0.0.0.0 --port 7860 ${CLI_ARGS}
diff --git a/services/lstein/mount.sh b/services/lstein/entrypoint.sh
old mode 100755
new mode 100644
similarity index 85%
rename from services/lstein/mount.sh
rename to services/lstein/entrypoint.sh
index 45bf13e..11dd2a6
--- a/services/lstein/mount.sh
+++ b/services/lstein/entrypoint.sh
@@ -7,9 +7,9 @@ declare -A MOUNTS
 # cache
 MOUNTS["/root/.cache"]=/data/.cache
 # ui specific
-MOUNTS["${PWD}/models/ldm/stable-diffusion-v1/model.ckpt"]=/data/StableDiffusion/model.ckpt
 MOUNTS["${PWD}/src/gfpgan/experiments/pretrained_models/GFPGANv1.4.pth"]=/data/GFPGAN/GFPGANv1.4.pth
 MOUNTS["${PWD}/ldm/invoke/restoration/codeformer/weights"]=/data/Codeformer
+MOUNTS["${PWD}/configs/models.yaml"]=/docker/models.yaml
 # hacks
 MOUNTS["/opt/conda/lib/python3.9/site-packages/facexlib/weights"]=/data/.cache
 MOUNTS["/opt/conda/lib/python3.9/site-packages/realesrgan/weights"]=/data/RealESRGAN
@@ -25,5 +25,7 @@ for to_path in "${!MOUNTS[@]}"; do
 done
 
 if "${PRELOAD}" == "true"; then
-  python3 -u scripts/preload_models.py
+  python3 -u scripts/preload_models.py --no-interactive
 fi
+
+exec "$@"
diff --git a/services/lstein/models.yaml b/services/lstein/models.yaml
new file mode 100644
index 0000000..4499cb2
--- /dev/null
+++ b/services/lstein/models.yaml
@@ -0,0 +1,15 @@
+# This file describes the alternative machine learning models
+# available to InvokeAI script.
+#
+# To add a new model, follow the examples below. Each
+# model requires a model config file, a weights file,
+# and the width and height of the images it
+# was trained on.
+stable-diffusion-1.5:
+  description: The newest Stable Diffusion version 1.5 weight file (4.27 GB)
+  weights: /data/StableDiffusion/v1-5-pruned-emaonly.ckpt
+  config: ./configs/stable-diffusion/v1-inference.yaml
+  width: 512
+  height: 512
+  vae: /data/VAE/vae-ft-mse-840000-ema-pruned.ckpt
+  default: true