From 249065f12775c0c931a68bd835c529a5281012a5 Mon Sep 17 00:00:00 2001
From: Kelvie Wong <kelvie@kelvie.ca>
Date: Sun, 28 May 2023 11:17:57 -0700
Subject: [PATCH] Add docker compose file

This makes it a lot easier to use locally on docker, and a lot easier to
move files in/out, as well as save notebooks.
---
 .gitignore         |  2 ++
 doc/SETUP.md       | 16 ++++++++++++++--
 docker-compose.yml | 19 +++++++++++++++++++
 docker/Dockerfile  |  4 +---
 docker/start.sh    |  6 ++++++
 5 files changed, 42 insertions(+), 5 deletions(-)
 create mode 100644 docker-compose.yml

diff --git a/.gitignore b/.gitignore
index bc8218c..9ff4db6 100644
--- a/.gitignore
+++ b/.gitignore
@@ -14,3 +14,5 @@
 .ssh_config
 *inference*.yaml
 .idea
+/.cache
+/models
diff --git a/doc/SETUP.md b/doc/SETUP.md
index 490372d..67b93a4 100644
--- a/doc/SETUP.md
+++ b/doc/SETUP.md
@@ -48,6 +48,18 @@ Double check your python version again after setup by running these two commands
 
 Again, this should show 3.10.x
 
-## Docker container
+## Local docker container
 
-`docker run -it -p 8888:8888 -p 6006:6006 --gpus all -e JUPYTER_PASSWORD=test1234 -t ghcr.io/victorchall/everydream2trainer:nightly`
+```sh
+    docker compose up
+```
+
+And you can either get a shell via:
+```sh
+    docker exec -it everydream2trainer-docker-everydream2trainer-1 /bin/bash
+```
+
+Or go to your browser and hit `http://localhost:8888`. The web password is
+`test1234` but you can change that in `docker-compose.yml`.
+
+Your current source directory will be moutned to the Jupyter notebook.
diff --git a/docker-compose.yml b/docker-compose.yml
new file mode 100644
index 0000000..8f4eca3
--- /dev/null
+++ b/docker-compose.yml
@@ -0,0 +1,19 @@
+version: '3.8'
+services:
+  everydream2trainer:
+    image: ghcr.io/victorchall/everydream2trainer:nightly
+    ports:
+      - "127.0.0.1:8888:8888"
+      - "127.0.0.1:6006:6006"
+    environment:
+      - JUPYTER_PASSWORD=test1234
+    volumes:
+      - .:/workspace/EveryDream2trainer
+      - ./.cache:/root/.cache
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
diff --git a/docker/Dockerfile b/docker/Dockerfile
index daadef6..3ec8ae1 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -87,10 +87,8 @@ ARG CACHEBUST=1
 RUN git clone https://github.com/victorchall/EveryDream2trainer
 
 WORKDIR /workspace/EveryDream2trainer
-RUN python utils/get_yamls.py && \
-    mkdir -p logs && mkdir -p input
 
 ADD welcome.txt /
 ADD start.sh /
 RUN chmod +x /start.sh
-CMD [ "/start.sh" ]
\ No newline at end of file
+CMD [ "/start.sh" ]
diff --git a/docker/start.sh b/docker/start.sh
index fbb5e3a..989f31d 100644
--- a/docker/start.sh
+++ b/docker/start.sh
@@ -2,6 +2,12 @@
 cat /welcome.txt
 export PYTHONUNBUFFERED=1
 
+if [[ ! -f "v2-inference-v.yaml" ]]; then
+    python utils/get_yamls.py
+fi
+
+mkdir -p logs input
+
 # RunPod SSH
 if [[ -v "PUBLIC_KEY" ]] && [[ ! -d "${HOME}/.ssh" ]]
 then