From 3880bb6e50291b89b216f4b40d0dc17ced61605b Mon Sep 17 00:00:00 2001
From: Chuck Tang <chuck.tang@databricks.com>
Date: Wed, 25 Sep 2024 15:16:32 -0700
Subject: [PATCH] commit change

---
 Dockerfile | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)
 create mode 100644 Dockerfile

diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 00000000..4900edf6
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1,21 @@
+# Copyright 2022 MosaicML Diffusion authors
+# SPDX-License-Identifier: Apache-2.0
+
+ARG BASE_IMAGE
+FROM $BASE_IMAGE
+
+ARG BRANCH_NAME
+ARG DEP_GROUPS
+
+ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 8.9 9.0"
+
+# Check for changes in setup.py.
+# If there are changes, the docker cache is invalidated and a fresh pip installation is triggered.
+ADD https://raw.githubusercontent.com/mosaicml/diffusion/$BRANCH_NAME/setup.py setup.py
+RUN rm setup.py
+
+# Install and uninstall diffusion to cache diffusion requirements
+RUN git clone -b $BRANCH_NAME https://github.com/mosaicml/diffusion.git
+RUN pip install --no-cache-dir "./diffusion${DEP_GROUPS}"
+RUN pip uninstall -y diffusion
+RUN rm -rf diffusion
\ No newline at end of file