RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Update docker install for blacwell

This commit is contained in:
Jaret Burkett 2025-06-01 19:06:53 -06:00
parent 12fa109910
commit e20a869dc1
2 changed files with 2 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
docker/Dockerfile
View File

@ -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

2
requirements.txt
View File

@ -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