From c3d571787d4701b3c4bb7c37538e2af670a2c417 Mon Sep 17 00:00:00 2001
From: YAN Wenkun <code@yanwk.fun>
Date: Mon, 26 Feb 2024 01:23:42 +0800
Subject: [PATCH] Tune strategy for installing xFormers

fix #20
---
 Dockerfile | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 3fc161c..48cb380 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -18,14 +18,23 @@ RUN --mount=type=cache,target=/var/cache/zypp \
         Mesa-libGL1 libgthread-2_0-0 \
     && rm /usr/lib64/python3.11/EXTERNALLY-MANAGED
 
-# Install xFormers (dev version, will install PyTorch as well)
 RUN --mount=type=cache,target=/root/.cache/pip \
     pip install --break-system-packages \
 	    --upgrade pip
 
+# Install xFormers (stable version, will specify PyTorch version)
+# and Torchvision + Torchaudio (will downgrade to match xFormers' PyTorch version)
 RUN --mount=type=cache,target=/root/.cache/pip \
     pip install --break-system-packages \
-        --pre -U xformers \
+        xformers torchvision torchaudio \
+        --index-url https://download.pytorch.org/whl/cu121 \
+        --extra-index-url https://pypi.org/simple
+
+# Upgrade xFormers to dev version
+# (While keeping most dependencies using stable version)
+RUN --mount=type=cache,target=/root/.cache/pip \
+    pip install --break-system-packages \
+        --pre --upgrade xformers \
         --index-url https://download.pytorch.org/whl/cu121 \
         --extra-index-url https://pypi.org/simple