From 69a18ee5ec5795f9b0f1937f8beb22fdb3747b4b Mon Sep 17 00:00:00 2001
From: Ashley Kleynhans <ashley.kleynhans@gmail.com>
Date: Wed, 13 Mar 2024 13:41:34 +0200
Subject: [PATCH] Use buildx bake to build the Docker image

---
 Dockerfile      | 22 ++++++++++++----------
 README.md       | 19 +++++++++++++++++++
 docker-bake.hcl | 29 +++++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 10 deletions(-)
 create mode 100644 docker-bake.hcl

diff --git a/Dockerfile b/Dockerfile
index b7c22b0..9e03058 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,10 +1,6 @@
 # Stage 1: Base
 FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04 as base
 
-ARG OOBABOOGA_COMMIT=1934cb61ef879815644277c01c7295acbae542d8
-ARG TORCH_VERSION=2.2.0
-ARG XFORMERS_VERSION=0.0.24
-
 SHELL ["/bin/bash", "-o", "pipefail", "-c"]
 ENV DEBIAN_FRONTEND=noninteractive \
     PYTHONUNBUFFERED=1 \
@@ -64,14 +60,18 @@ RUN ln -s /usr/bin/python3.10 /usr/bin/python
 FROM base as setup
 
 # Install Torch
+ARG INDEX_URL
+ARG TORCH_VERSION
+ARG XFORMERS_VERSION
+WORKDIR /
 RUN python3 -m venv /venv && \
     source /venv/bin/activate && \
-    pip3 install torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/cu121 && \
-    pip3 install xformers==${XFORMERS_VERSION} && \
+    pip3 install torch==${TORCH_VERSION} --index-url ${INDEX_URL} && \
+    pip3 install xformers==${XFORMERS_VERSION} --index-url ${INDEX_URL} && \
     deactivate
 
 # Clone the git repo of Text Generation Web UI and set version
-WORKDIR /
+ARG OOBABOOGA_COMMIT
 RUN git clone https://github.com/oobabooga/text-generation-webui && \
     cd /text-generation-webui && \
     git checkout ${OOBABOOGA_COMMIT}
@@ -98,7 +98,7 @@ RUN source /venv/bin/activate && \
 RUN curl https://rclone.org/install.sh | bash
 
 # Install runpodctl
-ARG RUNPODCTL_VERSION="v1.14.2"
+ARG RUNPODCTL_VERSION
 RUN wget "https://github.com/runpod/runpodctl/releases/download/${RUNPODCTL_VERSION}/runpodctl-linux-amd64" -O runpodctl && \
     chmod a+x runpodctl && \
     mv runpodctl /usr/local/bin
@@ -138,10 +138,12 @@ COPY fetch_model.py /text-generation-webui/
 COPY download_model.py /text-generation-webui/
 
 # Set template version
-ENV TEMPLATE_VERSION=1.13.1
+ARG RELEASE
+ENV TEMPLATE_VERSION=${RELEASE}
 
 # Set the venv path
-ENV VENV_PATH="/workspace/venvs/text-generation-webui"
+ARG VENV_PATH
+ENV VENV_PATH=${VENV_PATH}
 
 # Copy the scripts
 WORKDIR /
diff --git a/README.md b/README.md
index 535c58f..7a7f02a 100644
--- a/README.md
+++ b/README.md
@@ -34,6 +34,25 @@ You can use my custom [RunPod template](
 https://runpod.io/gsc?template=el5m58e1to&ref=2xxro4sy)
 to launch it on RunPod.
 
+## Building the Docker image
+
+> [!NOTE]
+> You will need to edit the `docker-bake.hcl` file and update `RELEASE`,
+> and `tags`.  You can obviously edit the other values too, but these
+> are the most important ones.
+
+```bash
+# Clone the repo
+git clone https://github.com/ashleykleynhans/text-generation-docker.git
+
+# Log in to Docker Hub
+docker login
+
+# Build the image, tag the image, and push the image to Docker Hub
+cd text-generation-docker
+docker buildx bake -f docker-bake.hcl --push
+```
+
 ## Running Locally
 
 ### Install Nvidia CUDA Driver
diff --git a/docker-bake.hcl b/docker-bake.hcl
new file mode 100644
index 0000000..9da6bc8
--- /dev/null
+++ b/docker-bake.hcl
@@ -0,0 +1,29 @@
+variable "USERNAME" {
+    default = "ashleykza"
+}
+
+variable "APP" {
+    default = "oobabooga"
+}
+
+variable "RELEASE" {
+    default = "1.13.2"
+}
+
+variable "CU_VERSION" {
+    default = "121"
+}
+
+target "default" {
+    dockerfile = "Dockerfile"
+    tags = ["${USERNAME}/${APP}:${RELEASE}"]
+    args = {
+        RELEASE = "${RELEASE}"
+        INDEX_URL = "https://download.pytorch.org/whl/cu${CU_VERSION}"
+        TORCH_VERSION = "2.2.0+cu${CU_VERSION}"
+        XFORMERS_VERSION = "0.0.24+cu${CU_VERSION}"
+        OOBABOOGA_COMMIT = "1934cb61ef879815644277c01c7295acbae542d8"
+        RUNPODCTL_VERSION = "v1.14.2"
+        VENV_PATH = "/workspace/venvs/text-generation-webuii"
+    }
+}