From 18c4a1c6c1f402e8a25078b4f55c015f8e3f5e54 Mon Sep 17 00:00:00 2001
From: andreasjansson <andreas@replicate.ai>
Date: Fri, 6 Sep 2024 21:49:58 +0200
Subject: [PATCH] Update to latest ai-toolkit, torch, and cuda

---
 ai-toolkit | 2 +-
 cog.yaml   | 8 +++++---
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/ai-toolkit b/ai-toolkit
index 169dbd2..279ee65 160000
--- a/ai-toolkit
+++ b/ai-toolkit
@@ -1 +1 @@
-Subproject commit 169dbd22ba3e3d6f1d8ed63c14b4ef3d8d1ab305
+Subproject commit 279ee65177a1ff7110fc05d6c2b7e082787ea8bc
diff --git a/cog.yaml b/cog.yaml
index 5f5c69e..eb3c084 100644
--- a/cog.yaml
+++ b/cog.yaml
@@ -3,12 +3,12 @@
 
 build:
   gpu: true
-  cuda: "12.1"
+  cuda: "12.4"
   python_version: "3.10.14"
   python_packages:
     # flux
-    - "torch==2.3.0"
-    - "torchvision==0.18.0"
+    - "torch==2.4.0"
+    - "torchvision==0.19.0"
     - "safetensors==0.4.4"
     - "diffusers==0.30.0"
     - "transformers==4.44.0"
@@ -56,6 +56,8 @@ build:
   run:
     - curl -o /usr/local/bin/pget -L "https://github.com/replicate/pget/releases/download/v0.8.2/pget_linux_x86_64" && chmod +x /usr/local/bin/pget
     - pip install "git+https://github.com/Gothos/diffusers.git@flux-inpaint"
+    - "pip install --force-reinstall --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu124"
+    #- "pip install --force-reinstall --pre torch==2.5.0.dev20240905+cu124 torchvision==0.20.0.dev20240905+cu124 --index-url https://download.pytorch.org/whl/nightly/cu124"
 
 # predict.py defines how predictions are run on your model
 predict: "predict.py:Predictor"