Bumped torch and cuda to support blackwell arch

README.md

@@ -114,7 +114,7 @@ cd ai-toolkit
 python3 -m venv venv
 source venv/bin/activate
 # 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
 ```
 
@@ -124,7 +124,7 @@ git clone https://github.com/ostris/ai-toolkit.git
 cd ai-toolkit
 python -m venv venv
 .\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
 ```
 

docker/Dockerfile

@@ -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"
 
@@ -46,13 +46,15 @@ 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.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
 ARG CACHEBUST=1234
+ARG GIT_COMMIT=main
 RUN echo "Cache bust: ${CACHEBUST}" && \
     git clone https://github.com/ostris/ai-toolkit.git && \
-    cd ai-toolkit
+    cd ai-toolkit && \
+    git checkout ${GIT_COMMIT}
 
 WORKDIR /app/ai-toolkit
 

requirements.txt

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

version.py

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