Merge pull request #129 from qslug/rp-torch2

Bump CUDA to 11.8 and Torch to 2.0 and xformers to 0.0.18

.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

@@ -1,6 +1,6 @@
 ###################
 # Builder Stage
-FROM nvidia/cuda:11.7.1-devel-ubuntu22.04 AS builder
+FROM nvidia/cuda:11.8.0-devel-ubuntu22.04 AS builder
 
 ARG DEBIAN_FRONTEND=noninteractive
 
@@ -23,19 +23,16 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
 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 && \
-    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
-
+    pip install -U -I torch==2.0.0+cu118 torchvision==0.15.1+cu118 --extra-index-url "https://download.pytorch.org/whl/cu118" && \
+    pip install -r requirements-build.txt && \
+    pip install --no-deps xformers==0.0.18
 
 ###################
 # Runtime Stage
-FROM nvidia/cuda:11.7.1-runtime-ubuntu22.04 as runtime
+FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04 as runtime
 
 # Use bash shell
 SHELL ["/bin/bash", "-o", "pipefail", "-c"]
@@ -74,13 +71,17 @@ RUN echo "source ${VIRTUAL_ENV}/bin/activate" >> /root/.bashrc
 # Workaround for:
 #   https://github.com/TimDettmers/bitsandbytes/issues/62
 #   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
+ENV LD_LIBRARY_PATH="/usr/local/cuda-11.8/targets/x86_64-linux/lib/"
+RUN ln /usr/local/cuda/targets/x86_64-linux/lib/libcudart.so.11.8.89 /usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudart.so
+RUN ln /usr/local/cuda/targets/x86_64-linux/lib/libnvrtc.so.11.8.89 /usr/local/cuda-11.8/targets/x86_64-linux/lib/libnvrtc.so
+ADD requirements-runtime.txt /
+RUN pip install --no-cache-dir -r requirements-runtime.txt
 
 WORKDIR /workspace
 RUN git clone https://github.com/victorchall/EveryDream2trainer
+
 WORKDIR /workspace/EveryDream2trainer
+# RUN git checkout torch2
 RUN python utils/get_yamls.py && \
     mkdir -p logs && mkdir -p input
 

docker/requirements-build.txt

@@ -0,0 +1,8 @@
+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.27.1

docker/requirements-runtime.txt

@@ -0,0 +1,13 @@
+aiohttp==3.8.4
+bitsandbytes==0.37.2
+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"

utils/fs_helpers.py

@@ -36,7 +36,11 @@ def walk_and_visit(path, visit_fn, context=None):
     files = []
     for name in names:
         fullname = os.path.join(path, name)
-        if os.path.isdir(fullname) and not str(name).startswith('.'):
+
+        if str(name).startswith('.'):
+            continue
+
+        if os.path.isdir(fullname):
             dirs.append(fullname)
         else:
             files.append(fullname)