From e20a869dc1d7bf93edffe3433d64126937c4544d Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Sun, 1 Jun 2025 19:06:53 -0600
Subject: [PATCH] Update docker install for blacwell

---
 docker/Dockerfile | 4 ++--
 requirements.txt  | 2 --
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 3093b067..69d1c974 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -49,7 +49,7 @@ WORKDIR /app
 RUN ln -s /usr/bin/python3 /usr/bin/python
 
 # install pytorch before cache bust to avoid redownloading pytorch
-RUN pip install --no-cache-dir torch==2.7.0 torchvision==0.22.0 --index-url https://download.pytorch.org/whl/cu128
+RUN pip install --pre --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu128
 
 # Fix cache busting by moving CACHEBUST to right before git clone
 ARG CACHEBUST=1234
@@ -63,7 +63,7 @@ WORKDIR /app/ai-toolkit
 
 # Install Python dependencies
 RUN pip install --no-cache-dir -r requirements.txt && \
-    pip install flash-attn --no-build-isolation  --no-cache-dir
+    pip install setuptools==69.5.1 --no-cache-dir
 
 # Build UI
 WORKDIR /app/ai-toolkit/ui
diff --git a/requirements.txt b/requirements.txt
index 7c51a320..ad7d00f2 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,3 @@
-torch==2.7.0
-torchvision==0.22.0
 torchao==0.10.0
 safetensors
 git+https://github.com/jaretburkett/easy_dwpose.git