From a9612122848d5458501f2bf3f87c58d1e65cf02e Mon Sep 17 00:00:00 2001
From: Augusto de la Torre <augustod@spotify.com>
Date: Wed, 29 Mar 2023 17:20:18 +0200
Subject: [PATCH] Split requirements between build and runtime

---
 .devcontainer/gpu/devcontainer.json    |  9 ++++++++-
 .devcontainer/no_gpu/devcontainer.json |  9 ++++++++-
 docker/Dockerfile                      |  7 ++++---
 docker/requirements-build.txt          |  9 +++++++++
 docker/requirements-runtime.txt        | 13 +++++++++++++
 docker/requirements.txt                | 20 --------------------
 docker/start.sh                        |  1 +
 7 files changed, 43 insertions(+), 25 deletions(-)
 create mode 100644 docker/requirements-build.txt
 create mode 100644 docker/requirements-runtime.txt
 delete mode 100644 docker/requirements.txt

diff --git a/.devcontainer/gpu/devcontainer.json b/.devcontainer/gpu/devcontainer.json
index 9e3f84c..59a816c 100644
--- a/.devcontainer/gpu/devcontainer.json
+++ b/.devcontainer/gpu/devcontainer.json
@@ -20,5 +20,12 @@
 
 	// Mimic RunPod/Vast setup
 	"workspaceMount": "source=${localWorkspaceFolder},target=/workspace/EveryDream2trainer,type=bind",
-	"workspaceFolder": "/workspace/EveryDream2trainer"
+	"workspaceFolder": "/workspace/EveryDream2trainer",
+	"customizations": {
+		"vscode": {
+			"extensions": [
+				"ms-python.python"
+			]
+		}
+	}
 }
\ No newline at end of file
diff --git a/.devcontainer/no_gpu/devcontainer.json b/.devcontainer/no_gpu/devcontainer.json
index 6b68898..6ede5bd 100644
--- a/.devcontainer/no_gpu/devcontainer.json
+++ b/.devcontainer/no_gpu/devcontainer.json
@@ -16,5 +16,12 @@
 	
 	// Mimic RunPod/Vast setup
 	"workspaceMount": "source=${localWorkspaceFolder},target=/workspace/EveryDream2trainer,type=bind",
-	"workspaceFolder": "/workspace/EveryDream2trainer"
+	"workspaceFolder": "/workspace/EveryDream2trainer",
+	"customizations": {
+		"vscode": {
+			"extensions": [
+				"ms-python.python"
+			]
+		}
+	}
 }
\ No newline at end of file
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 622c9a6..96f1002 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -24,11 +24,11 @@ ENV VIRTUAL_ENV=/workspace/venv
 ENV PATH="$VIRTUAL_ENV/bin:$PATH"
 
 
-ADD requirements.txt /build
+ADD requirements-build.txt /build
 RUN --mount=type=cache,target=/root/.cache/pip \
     python3 -m venv ${VIRTUAL_ENV} && \
     pip install -U -I torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url "https://download.pytorch.org/whl/cu117" && \
-    pip install -r requirements.txt && \
+    pip install -r requirements-build.txt && \
     export FORCE_CUDA=1 && export TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6" && export CUDA_VISIBLE_DEVICES=0 && \
     pip install --no-deps git+https://github.com/facebookresearch/xformers.git@e14dba4#egg=xformers
 
@@ -76,7 +76,8 @@ RUN echo "source ${VIRTUAL_ENV}/bin/activate" >> /root/.bashrc
 #   https://github.com/TimDettmers/bitsandbytes/issues/73
 ENV LD_LIBRARY_PATH="/usr/local/cuda-11.7/targets/x86_64-linux/lib"
 RUN ln /usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudart.so.11.0 /usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudart.so
-RUN pip install bitsandbytes==0.37.0
+ADD requirements-runtime.txt /
+RUN pip install --no-cache-dir -r requirements-runtime.txt
 
 WORKDIR /workspace
 RUN git clone https://github.com/victorchall/EveryDream2trainer
diff --git a/docker/requirements-build.txt b/docker/requirements-build.txt
new file mode 100644
index 0000000..83f3da4
--- /dev/null
+++ b/docker/requirements-build.txt
@@ -0,0 +1,9 @@
+diffusers[torch]>=0.14.0
+ninja
+numpy
+omegaconf==2.2.3
+protobuf==3.20.3
+pyre-extensions==0.0.23
+pytorch-lightning==1.9.2
+transformers==4.25.1
+triton>=2.0.0a2
\ No newline at end of file
diff --git a/docker/requirements-runtime.txt b/docker/requirements-runtime.txt
new file mode 100644
index 0000000..0e7a591
--- /dev/null
+++ b/docker/requirements-runtime.txt
@@ -0,0 +1,13 @@
+aiohttp==3.8.4
+bitsandbytes==0.37.0
+colorama==0.4.6
+ftfy==6.1.1
+ipyevents
+ipywidgets
+jupyter-archive
+jupyterlab
+piexif==1.1.3
+pyfakefs
+pynvml==11.5.0
+tensorboard==2.12.0
+wandb
diff --git a/docker/requirements.txt b/docker/requirements.txt
deleted file mode 100644
index 5fc1856..0000000
--- a/docker/requirements.txt
+++ /dev/null
@@ -1,20 +0,0 @@
-aiohttp==3.8.4
-colorama==0.4.6
-diffusers[torch]>=0.14.0
-ftfy==6.1.1
-ipyevents
-ipywidgets
-jupyter-archive
-jupyterlab
-ninja
-omegaconf==2.2.3
-piexif==1.1.3
-protobuf==3.20.3
-pyfakefs
-pynvml==11.5.0
-pyre-extensions==0.0.30
-pytorch-lightning==1.9.2
-tensorboard==2.12.0
-transformers==4.25.1
-triton>=2.0.0a2
-wandb
\ No newline at end of file
diff --git a/docker/start.sh b/docker/start.sh
index 67bf308..a6e4e90 100644
--- a/docker/start.sh
+++ b/docker/start.sh
@@ -18,6 +18,7 @@ if [[ $JUPYTER_PASSWORD ]]
 then
     tensorboard --logdir /workspace/EveryDream2trainer/logs --host 0.0.0.0 &
     jupyter nbextension enable --py widgetsnbextension
+    jupyter labextension disable "@jupyterlab/apputils-extension:announcements"
     jupyter lab --allow-root --no-browser --port=8888 --ip=* --ServerApp.terminado_settings='{"shell_command":["/bin/bash"]}' --ServerApp.token=$JUPYTER_PASSWORD --ServerApp.allow_origin=* --ServerApp.preferred_dir=/workspace/EveryDream2trainer
 else
     echo "Container Started"