From ed36edd85b886623377beb5f50c5dd7e3b3eb89a Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Tue, 23 Dec 2025 19:44:20 -0700
Subject: [PATCH] Bumped docker OS, python version, torch version

---
 docker/Dockerfile | 14 ++++++++------
 requirements.txt  |  3 ++-
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index a4f68f44..44d4d9db 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.8.1-devel-ubuntu22.04
+FROM nvidia/cuda:12.8.1-devel-ubuntu24.04
 
 LABEL authors="jaret"
 
@@ -15,7 +15,7 @@ RUN apt-get update && apt-get install --no-install-recommends -y \
     build-essential \
     cmake \
     wget \
-    python3.10 \
+    python3.12 \
     python3-pip \
     python3-dev \
     python3-setuptools \
@@ -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 --pre --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
+RUN pip install --no-cache-dir torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cu128 --break-system-packages
 
 # Fix cache busting by moving CACHEBUST to right before git clone
 ARG CACHEBUST=1234
@@ -61,10 +61,12 @@ RUN echo "Cache bust: ${CACHEBUST}" && \
 
 WORKDIR /app/ai-toolkit
 
+# copy requirements.txt to leverage docker cache
+COPY requirements.txt /app/ai-toolkit/requirements.txt
+
 # Install Python dependencies
-RUN pip install --no-cache-dir -r requirements.txt && \
-    pip install --pre --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128 --force && \
-    pip install setuptools==69.5.1 --no-cache-dir
+RUN pip install --no-cache-dir --break-system-packages -r requirements.txt && \
+    pip install setuptools==69.5.1 --no-cache-dir --break-system-packages
 
 # Build UI
 WORKDIR /app/ai-toolkit/ui
diff --git a/requirements.txt b/requirements.txt
index 577dde87..a4a7c5e0 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -35,4 +35,5 @@ python-slugify
 opencv-python
 pytorch-wavelets==1.3.0
 matplotlib==3.10.1
-setuptools==69.5.1
\ No newline at end of file
+setuptools==69.5.1
+scipy==1.12.0
\ No newline at end of file