allow prerelease xformers

neggles · Feb 26, 2024 · d72fe44 · d72fe44
1 parent c3867c3
commit d72fe44
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/docker/base/Dockerfile b/docker/base/Dockerfile
@@ -167,8 +167,8 @@ FROM base as xformers-binary
 ARG XFORMERS_PACKAGE="xformers"
 ARG XFORMERS_PIP_ARGS=""
 RUN --mount=type=cache,target=/root/.cache/pip,sharing=locked \
-    python -m pip install ${XFORMERS_PIP_ARGS} "${XFORMERS_PACKAGE}"
-
+    python -m pip install ${XFORMERS_PIP_ARGS} "${XFORMERS_PACKAGE}" \
+    || python -m pip install --pre ${XFORMERS_PIP_ARGS} "${XFORMERS_PACKAGE}"
 
 # or if xformers has to be compiled; flash attention 2 makes this impossible
 # on a public github actions runner without limiting yourself to only a handful of architectures