From 2bab2c7592b8df13a3e333ee51eb24f7fd086265 Mon Sep 17 00:00:00 2001
From: Rob Ballantyne <rob@robballantyne.co.uk>
Date: Mon, 18 Nov 2024 18:10:00 +0000
Subject: [PATCH] Build 5.4.1

---
 .github/workflows/docker-build.yml            |  2 +-
 .../opt/ai-dock/bin/build/layer99/init.sh     | 22 +++++++++++++------
 docker-compose.yaml                           | 18 ++++++++++-----
 3 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml
index 7ccb1ee..de4e125 100644
--- a/.github/workflows/docker-build.yml
+++ b/.github/workflows/docker-build.yml
@@ -88,7 +88,7 @@ jobs:
       fail-fast: false
       matrix:
         build:
-          - {latest: "false", invokeai: "4.2.7", python: "3.10", pytorch: "2.2.2", cuda: "12.1.1-base"}
+          - {latest: "false", invokeai: "5.4.1", python: "3.10", pytorch: "2.4.1", cuda: "12.1.1-base"}
           
     steps:
       -
diff --git a/build/COPY_ROOT_99/opt/ai-dock/bin/build/layer99/init.sh b/build/COPY_ROOT_99/opt/ai-dock/bin/build/layer99/init.sh
index 756fbd7..cbaa18c 100755
--- a/build/COPY_ROOT_99/opt/ai-dock/bin/build/layer99/init.sh
+++ b/build/COPY_ROOT_99/opt/ai-dock/bin/build/layer99/init.sh
@@ -80,25 +80,33 @@ function build_extra_start() {
 
     # Invoke has no exit/CI run mode so run it and wait until it's fuly initialised
     source "$INVOKEAI_VENV/bin/activate"
-    LD_PRELOAD=libtcmalloc.so invokeai-web 2>&1 | tee /tmp/invoke-ci.log &
+    LD_PRELOAD=libtcmalloc.so invokeai-web 2>&1 | tee /tmp/invoke-ci.log & 
+    INVOKE_PID=$!
+
     wait_max=30
     wait_current=0
     init_string="Uvicorn running on"
 
-    # Until loop to continuously check if the string is found or maximum wait time is reached
     while [ $wait_current -lt $wait_max ]; do
         if grep -qi "$init_string" /tmp/invoke-ci.log; then
             echo "InvokeAI initialization complete."
+            kill -9 $INVOKE_PID
             break
-        else
-            echo "Waiting for InvokeAI initialization to complete..."
-            sleep 1
-            wait_current=$((wait_current + 1))
         fi
+        echo "Waiting for InvokeAI initialization..."
+        sleep 1
+        wait_current=$((wait_current + 1))
     done
 
-    pkill invokeai-web
+    if [ $wait_current -eq $wait_max ]; then
+        echo "Timeout waiting for InvokeAI"
+        kill -9 $INVOKE_PID
+        exit 1
+    fi
+    ps aux
+    echo "Should be exiting about now"
     deactivate
+    echo "Still waiting?"
 }
 
 function build_extra_get_pip_packages() {
diff --git a/docker-compose.yaml b/docker-compose.yaml
index 6b55024..e761e63 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -7,14 +7,22 @@ services:
       context: ./build
       args:
         PYTHON_VERSION: ${PYTHON_VERSION:-3.10}
-        PYTORCH_VERSION: ${PYTORCH_VERSION:-2.2.2}
-        INVOKEAI_VERSION: ${INVOKEAI_VERSION:-4.2.4}
-        IMAGE_BASE: ${IMAGE_BASE:-ghcr.io/ai-dock/python:3.10-v2-cuda-11.8.0-base-22.04}
+        PYTORCH_VERSION: ${PYTORCH_VERSION:-2.4.1}
+        INVOKEAI_VERSION: ${INVOKEAI_VERSION:-5.4.1}
+        IMAGE_BASE: ${IMAGE_BASE:-ghcr.io/ai-dock/python:3.10-v2-cuda-12.1.1-base-22.04}
       tags:
-        - "ghcr.io/ai-dock/invokeai:${IMAGE_TAG:-v2-cuda-11.8.0-base-22.04}"
+        - "ghcr.io/ai-dock/invokeai:${IMAGE_TAG:-v2-cuda-12.1.1-base-22.04}"
         
-    image: ghcr.io/ai-dock/invokeai:${IMAGE_TAG:-v2-cuda-11.8.0-base-22.04}
+    image: ghcr.io/ai-dock/invokeai:${IMAGE_TAG:-v2-cuda-12.1.1-base-22.04}
     
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
+
     devices:
       - "/dev/dri:/dev/dri"
       # For AMD GPU