Adjustments

H3mul · Dec 7, 2023 · 2f217ac · 2f217ac
1 parent 7b07deb
commit 2f217ac
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 16 deletions.
diff --git a/scripts/track-training-pids.sh → scripts/pid-tracker.sh b/scripts/track-training-pids.sh → scripts/pid-tracker.sh
@@ -2,20 +2,31 @@
 
 set -eu
 
+trap clean_up ERR EXIT SIGINT SIGTERM
+
+SCRIPT_PID_FILE="/app/pid/pid-tracker.pid"
+
+clean_up() {
+    trap - ERR EXIT SIGINT SIGTERM
+    # Remove temporary files/directories, log files or rollback changes.
+    rm ${SCRIPT_PID_FILE}
+}
+
+
 if [ -f /app/pid/track-training-pids.pid ];then
     echo "Killing running instance of this script..."
-    rkill $(cat /app/pid/track-training-pids.pid) || true
+    rkill $(cat ${SCRIPT_PID_FILE}) || true
 fi
 
-echo $$ > /app/pid/track-training-pids.pid
+echo $$ > ${SCRIPT_PID_FILE}
 
 # This script runs in the background and tracks common long-running tasks
 # which we might want to quit and/or notify the completion of.
 
-mkdir -p /app/pid/training-pids
+mkdir -p /app/pid/tracked-pids
 
 echo "Clearing past tracked PIDs..."
-rm /app/pid/training-pids/*
+rm /app/pid/tracked-pids/*
 
 echo "Listening for training PIDs..."
 while true; do
@@ -27,15 +38,15 @@ while true; do
         [ -z "${model_name}" ] && continue
 
         # 3. Add a file to /app/pid with PID and info string
-        pid_file="/app/pid/training-pids/${model_name}.pid"
+        pid_file="/app/pid/tracked-pids/${model_name}.pid"
         if [ ! -f ${pid_file} ]; then
             echo "Found a new model being trained: ${model_name}, tracking..."
             echo "${pid}" > ${pid_file}
         fi
     done
 
     # 4. Check all currently tracked pids to see if any exited, if so quit and/or notify
-    for pid_file in /app/pid/training-pids/*.pid; do
+    for pid_file in /app/pid/tracked-pids/*.pid; do
         model_name=$(basename ${pid_file} .pid)
         pid=$(cat ${pid_file})
 

diff --git a/scripts/pre_start.sh b/scripts/pre_start.sh
@@ -12,8 +12,8 @@ echo "Setting the app install root to ${INSTALL_ROOT}"
 mkdir -p ${INSTALL_ROOT}
 
 echo "Starting app provision in parallel..."
-# provision_invokeai.sh &
-# provision_kohya.sh &
+provision_invokeai.sh &
+provision_kohya.sh &
 provision_a1111.sh &
 wait
 echo "All app installs complete"

diff --git a/scripts/provision_a1111.sh b/scripts/provision_a1111.sh
@@ -30,14 +30,15 @@ echo "Downloading SDXL styles..."
 aria2c https://raw.githubusercontent.com/Douleb/SDXL-750-Styles-GPT4-/main/styles.csv -o ./styles.csv
 
 echo "Installing plugins..."
-git clone https://github.com/d8ahazard/sd_dreambooth_extension.git extensions/sd_dreambooth_extension
-git clone --depth=1 https://github.com/deforum-art/sd-webui-deforum.git extensions/deforum
-git clone --depth=1 https://github.com/Mikubill/sd-webui-controlnet.git extensions/sd-webui-controlnet
-git clone --depth=1 https://github.com/ashleykleynhans/a1111-sd-webui-locon.git extensions/a1111-sd-webui-locon
-git clone --depth=1 https://github.com/Gourieff/sd-webui-reactor.git extensions/sd-webui-reactor
-git clone --depth=1 https://github.com/zanllp/sd-webui-infinite-image-browsing.git extensions/infinite-image-browsing
-git clone --depth=1 https://github.com/Bing-su/adetailer.git extensions/adetailer
+[ -d extensions/sd_dreambooth_extension ] || git clone https://github.com/d8ahazard/sd_dreambooth_extension.git extensions/sd_dreambooth_extension
+[ -d extensions/deforum ] || git clone --depth=1 https://github.com/deforum-art/sd-webui-deforum.git extensions/deforum
+[ -d extensions/sd-webui-controlnet ] || git clone --depth=1 https://github.com/Mikubill/sd-webui-controlnet.git extensions/sd-webui-controlnet
+[ -d extensions/a1111-sd-webui-locon ] || git clone --depth=1 https://github.com/ashleykleynhans/a1111-sd-webui-locon.git extensions/a1111-sd-webui-locon
+[ -d extensions/sd-webui-reactor ] || git clone --depth=1 https://github.com/Gourieff/sd-webui-reactor.git extensions/sd-webui-reactor
+[ -d extensions/infinite-image-browsing ] || git clone --depth=1 https://github.com/zanllp/sd-webui-infinite-image-browsing.git extensions/infinite-image-browsing
+[ -d extensions/adetailer ] || git clone --depth=1 https://github.com/Bing-su/adetailer.git extensions/adetailer
 
+echo "Installing plugin dependencies..."
 cd ${A1111_ROOT}/extensions/deforum
 pip3 install -r requirements.txt
 cd ${A1111_ROOT}/extensions/sd-webui-controlnet
@@ -57,6 +58,8 @@ git reset ${DREAMBOOTH_COMMIT} --hard
 cp /app/config/a1111/requirements_dreambooth.txt ./requirements.txt
 pip3 install -r requirements.txt
 
+echo "Finished installing plugins..."
+
 cp /app/config/a1111/{webui-user.sh,config.json,ui-config.json} ./
 
 pip3 cache purge

diff --git a/scripts/send-notification.sh b/scripts/send-notification.sh
@@ -17,4 +17,4 @@ fi
 curl -u ${PUSHBULLET_API_TOKEN}: \
 	-X POST https://api.pushbullet.com/v2/pushes \
 	--header 'Content-Type: application/json' \
-	--data-binary "{\"type\":\"note\",\"title\":\"sd-ultimate\",\"body\":\"${MESSAGE}\"}"
+	--data-binary "{\"type\":\"note\",\"title\":\"sd-ultimate\",\"body\":\"${MESSAGE}\"}" >/dev/null 2>&1
diff --git a/scripts/start.sh b/scripts/start.sh
@@ -55,6 +55,10 @@ start_vsserver () {
     fi
 }
 
+start_pid_tracker() {
+    nohup pid-tracker.sh > /workspace/logs/pid-tracker.log 2>&1 &
+}
+
 # ---------------------------------------------------------------------------- #
 #                               Main Program                                   #
 # ---------------------------------------------------------------------------- #
@@ -81,6 +85,7 @@ if [ "${DISABLE_AUTOLAUNCH}" != true ]; then
 fi
 
 start_vsserver
+start_pid_tracker
 
 echo "Container is READY!"