Add docker build action

And also a VS Code dev container

.devcontainer/devcontainer.json

@@ -0,0 +1,13 @@
+{
+    "name": "EveryDream2 Dev Container",
+	"dockerFile": "../docker/Dockerfile",
+	"postStartCommand": "/start.sh",
+
+	"containerEnv": {
+		"LOCAL_DEV": "1"
+	},
+
+	// Mimic RunPod/Vast setup
+	"workspaceMount": "source=${localWorkspaceFolder},target=/workspace/EveryDream2trainer,type=bind",
+	"workspaceFolder": "/workspace/EveryDream2trainer"
+}

.github/workflows/docker-image.yml

@@ -0,0 +1,31 @@
+name: Docker Image CI
+
+on:
+  push:
+    branches: [ "*" ]
+  pull_request:
+    branches: [ "main" ]
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v3
+
+    - name: Set up Docker Buildx
+      uses: docker/setup-buildx-action@v2
+
+    # Build Docker image with Buildx
+    # https://github.com/docker/build-push-action
+    - name: Build and push
+      id: build-and-push
+      uses: docker/build-push-action@v4
+      with:
+        context: .
+        push: false
+        file: docker/Dockerfile
+        cache-from: type=gha
+        cache-to: type=gha,mode=max

.github/workflows/docker-publish.yml

@@ -0,0 +1,73 @@
+name: Docker
+
+# This workflow uses actions that are not certified by GitHub.
+# They are provided by a third-party and are governed by
+# separate terms of service, privacy policy, and support
+# documentation.
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: '17 6 * * *'
+  push:
+    branches: [ "main" ]
+    # Publish semver tags as releases.
+    tags: [ 'v*.*.*' ]
+  pull_request:
+    branches: [ "main" ]
+
+env:
+  # Use docker.io for Docker Hub if empty
+  REGISTRY: ghcr.io
+  # github.repository as <account>/<repo>
+  IMAGE_NAME: ${{ github.repository }}
+
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v3
+
+#      - name: Free Disk Space (Ubuntu)
+#        uses: jlumbroso/free-disk-space@main
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+
+      # Login against a Docker registry except on PR
+      # https://github.com/docker/login-action
+      - name: Docker login
+        uses: docker/login-action@v2
+        with:
+          registry: ${{ env.REGISTRY }}
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      # Extract metadata (tags, labels) for Docker
+      # https://github.com/docker/metadata-action
+      - name: Docker meta
+        id: meta
+        uses: docker/metadata-action@v4
+        with:
+          images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
+
+      # Build and push Docker image with Buildx
+      # https://github.com/docker/build-push-action
+      - name: Build and push
+        id: build-and-push
+        uses: docker/build-push-action@v4
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.meta.outputs.tags }}
+          labels: ${{ steps.meta.outputs.labels }}
+          file: docker/Dockerfile
+          cache-from: type=gha
+          cache-to: type=gha,mode=max

.gitignore

@@ -10,3 +10,5 @@
 **/.ipynb_checkpoints
 /wandb/**
 /mycfgs/**
+/.vscode/**
+.ssh_config

docker/Dockerfile

@@ -0,0 +1,86 @@
+###################
+# Builder Stage
+FROM nvidia/cuda:11.7.1-devel-ubuntu22.04 AS builder
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+# Don't write .pyc bytecode
+ENV PYTHONDONTWRITEBYTECODE=1
+
+# Create workspace working directory
+RUN mkdir /build
+WORKDIR /build
+
+RUN rm -f /etc/apt/apt.conf.d/docker-clean; echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' > /etc/apt/apt.conf.d/keep-cache
+RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
+    --mount=type=cache,target=/var/lib/apt,sharing=locked \
+    apt update && apt-get install -y \
+        git wget build-essential \
+        python3-venv python3-pip \
+        gnupg ca-certificates \
+    && update-ca-certificates
+
+ENV VIRTUAL_ENV=/workspace/venv
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+
+
+ADD docker/requirements.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 --pre --no-deps xformers==0.0.17.dev451
+#    In case of emergency, build xformers from scratch
+#    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@48a77cc#egg=xformers
+
+
+###################
+# Runtime Stage
+FROM nvidia/cuda:11.7.1-runtime-ubuntu22.04 as runtime
+
+# Use bash shell
+SHELL ["/bin/bash", "-o", "pipefail", "-c"]
+ENV DEBIAN_FRONTEND noninteractive\
+    SHELL=/bin/bash
+
+# Python logs go strait to stdout/stderr w/o buffering
+ENV PYTHONUNBUFFERED=1
+
+# Don't write .pyc bytecode
+ENV PYTHONDONTWRITEBYTECODE=1
+
+RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
+    --mount=type=cache,target=/var/lib/apt,sharing=locked \
+    apt update && apt install -y --no-install-recommends \
+        wget bash curl git git-lfs vim tmux \
+        build-essential lsb-release \
+        python3-pip python3-venv \
+        openssh-server \
+        gnupg ca-certificates && \
+    update-ca-certificates && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/* && \
+    echo "en_US.UTF-8 UTF-8" > /etc/locale.gen
+
+# Install runpodctl
+RUN wget https://github.com/runpod/runpodctl/releases/download/v1.9.0/runpodctl-linux-amd -O runpodctl && \
+    chmod a+x runpodctl && \
+    mv runpodctl /usr/local/bin
+
+ENV VIRTUAL_ENV=/workspace/venv
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+COPY --from=builder ${VIRTUAL_ENV} ${VIRTUAL_ENV}
+
+# 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
+
+WORKDIR /workspace
+
+ADD docker/welcome.txt /
+ADD docker/start.sh /
+RUN chmod +x /start.sh
+CMD [ "/start.sh" ]

docker/requirements.txt

@@ -0,0 +1,19 @@
+aiohttp==3.8.4
+colorama==0.4.6
+diffusers[torch]>=0.13.0
+ftfy==6.1.1
+ipyevents
+ipywidgets
+jupyter-archive
+jupyterlab
+ninja
+omegaconf==2.2.3
+piexif==1.1.3
+protobuf==3.20.3
+pynvml==11.5.0
+pyre-extensions==0.0.30
+pytorch-lightning==1.9.2
+tensorboard==2.11.0
+transformers==4.25.1
+triton>=2.0.0a2
+wandb

docker/start.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+cat /welcome.txt
+export PYTHONUNBUFFERED=1
+
+echo "source /workspace/venv/bin/activate" >> ~/.bashrc
+source ~/.bashrc
+
+# Workaround for:
+#   https://github.com/TimDettmers/bitsandbytes/issues/62
+#   https://github.com/TimDettmers/bitsandbytes/issues/73
+pip install bitsandbytes==0.37.0
+
+function clone_pull {
+  DIRECTORY=$(basename "$1" .git)
+  if [ -d "$DIRECTORY" ]; then
+    cd "$DIRECTORY"
+    git pull
+    cd ../
+  else
+    git clone "$1"
+  fi
+}
+
+
+# VSCode Dev Container
+if [[ $LOCAL_DEV ]]
+then
+  echo "Running in dev container, skipping git pull"
+else
+  clone_pull https://github.com/victorchall/EveryDream2trainer
+fi
+cd /workspace/EveryDream2trainer
+python utils/get_yamls.py
+mkdir -p /workspace/EveryDream2trainer/logs
+mkdir -p /workspace/EveryDream2trainer/input
+
+# RunPod SSH
+if [[ -v "PUBLIC_KEY" ]] && [[ ! -d "${HOME}/.ssh" ]]
+then
+    pushd $HOME
+    mkdir -p .ssh
+    echo ${PUBLIC_KEY} > .ssh/authorized_keys
+    chmod -R 700 .ssh
+    popd
+    service ssh start
+fi
+
+# RunPod JupyterLab
+if [[ $JUPYTER_PASSWORD ]]
+then
+    tensorboard --logdir /workspace/EveryDream2trainer/logs --host 0.0.0.0 &
+    jupyter nbextension enable --py widgetsnbextension
+    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"
+    sleep infinity
+fi

docker/welcome.txt

@@ -0,0 +1,14 @@
+
+
+███████╗██╗   ██╗███████╗██████╗ ██╗   ██╗██████╗ ██████╗ ███████╗ █████╗ ███╗   ███╗
+██╔════╝██║   ██║██╔════╝██╔══██╗╚██╗ ██╔╝██╔══██╗██╔══██╗██╔════╝██╔══██╗████╗ ████║
+█████╗  ██║   ██║█████╗  ██████╔╝ ╚████╔╝ ██║  ██║██████╔╝█████╗  ███████║██╔████╔██║
+██╔══╝  ╚██╗ ██╔╝██╔══╝  ██╔══██╗  ╚██╔╝  ██║  ██║██╔══██╗██╔══╝  ██╔══██║██║╚██╔╝██║
+███████╗ ╚████╔╝ ███████╗██║  ██║   ██║   ██████╔╝██║  ██║███████╗██║  ██║██║ ╚═╝ ██║
+╚══════╝  ╚═══╝  ╚══════╝╚═╝  ╚═╝   ╚═╝   ╚═════╝ ╚═╝  ╚═╝╚══════╝╚═╝  ╚═╝╚═╝     ╚═╝ 2
+
+
+Support:
+--------
+
+Repo: https://github.com/victorchall/EveryDream2trainer