Split requirements between build and runtime

.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"
+			]
+		}
+	}
 }

.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"
+			]
+		}
+	}
 }

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

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

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

docker/requirements.txt

@@ -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

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"