Use buildx bake to build the Docker image

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 /

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

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