github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Bumped torch and cuda to support blackwell arch

Browse Source
This commit is contained in:
Jaret Burkett
2025-05-09 11:17:58 -06:00
parent 66a41e49d9
commit 64130189ce
4 changed files with 11 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
README.md
View File

@@ -114,7 +114,7 @@ cd ai-toolkit
python3 -m venv venv python3 -m venv venv
source venv/bin/activate source venv/bin/activate
# install torch first # install torch first
pip3 install --no-cache-dir torch==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu126 pip3 install --no-cache-dir torch==2.7.0 torchvision==0.22.0 --index-url https://download.pytorch.org/whl/cu126
pip3 install -r requirements.txt pip3 install -r requirements.txt
``` ```
@@ -124,7 +124,7 @@ git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit cd ai-toolkit
python -m venv venv python -m venv venv
.\venv\Scripts\activate .\venv\Scripts\activate
pip install --no-cache-dir torch==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu126 pip install --no-cache-dir torch==2.7.0 torchvision==0.22.0 --index-url https://download.pytorch.org/whl/cu126
pip install -r requirements.txt pip install -r requirements.txt
``` ```

8
docker/Dockerfile
View File

@@ -1,4 +1,4 @@
FROM nvidia/cuda:12.6.3-devel-ubuntu22.04 FROM nvidia/cuda:12.8.1-devel-ubuntu22.04
LABEL authors="jaret" LABEL authors="jaret"
@@ -46,13 +46,15 @@ WORKDIR /app
RUN ln -s /usr/bin/python3 /usr/bin/python RUN ln -s /usr/bin/python3 /usr/bin/python
# install pytorch before cache bust to avoid redownloading pytorch # install pytorch before cache bust to avoid redownloading pytorch
RUN pip install --no-cache-dir torch==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu126 RUN pip install --no-cache-dir torch==2.7.0 torchvision==0.22.0 --index-url https://download.pytorch.org/whl/cu128
# Fix cache busting by moving CACHEBUST to right before git clone # Fix cache busting by moving CACHEBUST to right before git clone
ARG CACHEBUST=1234 ARG CACHEBUST=1234
ARG GIT_COMMIT=main
RUN echo "Cache bust: ${CACHEBUST}" && \ RUN echo "Cache bust: ${CACHEBUST}" && \
git clone https://github.com/ostris/ai-toolkit.git && \ git clone https://github.com/ostris/ai-toolkit.git && \
cd ai-toolkit cd ai-toolkit && \
git checkout ${GIT_COMMIT}
WORKDIR /app/ai-toolkit WORKDIR /app/ai-toolkit

6
requirements.txt
View File

@@ -1,6 +1,6 @@
torch==2.6.0 torch==2.7.0
torchvision==0.21.0 torchvision==0.22.0
torchao==0.9.0 torchao==0.10.0
safetensors safetensors
git+https://github.com/jaretburkett/easy_dwpose.git git+https://github.com/jaretburkett/easy_dwpose.git
git+https://github.com/huggingface/diffusers@363d1ab7e24c5ed6c190abb00df66d9edb74383b git+https://github.com/huggingface/diffusers@363d1ab7e24c5ed6c190abb00df66d9edb74383b

2
version.py
View File

@@ -1 +1 @@
VERSION = "0.2.8" VERSION = "0.2.9"