Merge branch 'dev' into release_candidate

1 files changed, 44 insertions, 19 deletions

diff --git a/webui.sh b/webui.sh
index ab52ac3b..5c8d977c 100755
--- a/webui.sh
+++ b/webui.sh
@@ -112,9 +112,24 @@ then
 fi
 
 # Check prerequisites
-gpu_info=$(lspci 2>/dev/null | grep VGA)
+gpu_info=$(lspci 2>/dev/null | grep -E "VGA|Display")
 case "$gpu_info" in
-    *"Navi 1"*|*"Navi 2"*) export HSA_OVERRIDE_GFX_VERSION=10.3.0
+    *"Navi 1"*)
+        export HSA_OVERRIDE_GFX_VERSION=10.3.0
+        if [[ -z "${TORCH_COMMAND}" ]]
+        then
+            pyv="$(${python_cmd} -c 'import sys; print(".".join(map(str, sys.version_info[0:2])))')"
+            if [[ $(bc <<< "$pyv <= 3.10") -eq 1 ]] 
+            then
+                # Navi users will still use torch 1.13 because 2.0 does not seem to work.
+                export TORCH_COMMAND="pip install torch==1.13.1+rocm5.2 torchvision==0.14.1+rocm5.2 --index-url https://download.pytorch.org/whl/rocm5.2"
+            else
+                printf "\e[1m\e[31mERROR: RX 5000 series GPUs must be using at max python 3.10, aborting...\e[0m"
+                exit 1
+            fi
+        fi
+    ;;
+    *"Navi 2"*) export HSA_OVERRIDE_GFX_VERSION=10.3.0
     ;;
     *"Renoir"*) export HSA_OVERRIDE_GFX_VERSION=9.0.0
         printf "\n%s\n" "${delimiter}"
@@ -124,9 +139,12 @@ case "$gpu_info" in
     *)
     ;;
 esac
-if echo "$gpu_info" | grep -q "AMD" && [[ -z "${TORCH_COMMAND}" ]]
+if ! echo "$gpu_info" | grep -q "NVIDIA";
 then
-    export TORCH_COMMAND="pip install torch==2.0.1+rocm5.4.2 torchvision==0.15.2+rocm5.4.2 --index-url https://download.pytorch.org/whl/rocm5.4.2"
+    if echo "$gpu_info" | grep -q "AMD" && [[ -z "${TORCH_COMMAND}" ]]
+    then
+        export TORCH_COMMAND="pip install torch==2.0.1+rocm5.4.2 torchvision==0.15.2+rocm5.4.2 --index-url https://download.pytorch.org/whl/rocm5.4.2"
+    fi
 fi
 
 for preq in "${GIT}" "${python_cmd}"
@@ -190,7 +208,7 @@ fi
 # Try using TCMalloc on Linux
 prepare_tcmalloc() {
     if [[ "${OSTYPE}" == "linux"* ]] && [[ -z "${NO_TCMALLOC}" ]] && [[ -z "${LD_PRELOAD}" ]]; then
-        TCMALLOC="$(ldconfig -p | grep -Po "libtcmalloc.so.\d" | head -n 1)"
+        TCMALLOC="$(PATH=/usr/sbin:$PATH ldconfig -p | grep -Po "libtcmalloc(_minimal|)\.so\.\d" | head -n 1)"
         if [[ ! -z "${TCMALLOC}" ]]; then
             echo "Using TCMalloc: ${TCMALLOC}"
             export LD_PRELOAD="${TCMALLOC}"
@@ -200,17 +218,24 @@ prepare_tcmalloc() {
     fi
 }
 
-if [[ ! -z "${ACCELERATE}" ]] && [ ${ACCELERATE}="True" ] && [ -x "$(command -v accelerate)" ]
-then
-    printf "\n%s\n" "${delimiter}"
-    printf "Accelerating launch.py..."
-    printf "\n%s\n" "${delimiter}"
-    prepare_tcmalloc
-    exec accelerate launch --num_cpu_threads_per_process=6 "${LAUNCH_SCRIPT}" "$@"
-else
-    printf "\n%s\n" "${delimiter}"
-    printf "Launching launch.py..."
-    printf "\n%s\n" "${delimiter}"
-    prepare_tcmalloc
-    exec "${python_cmd}" "${LAUNCH_SCRIPT}" "$@"
-fi
+KEEP_GOING=1
+export SD_WEBUI_RESTART=tmp/restart
+while [[ "$KEEP_GOING" -eq "1" ]]; do
+    if [[ ! -z "${ACCELERATE}" ]] && [ ${ACCELERATE}="True" ] && [ -x "$(command -v accelerate)" ]; then
+        printf "\n%s\n" "${delimiter}"
+        printf "Accelerating launch.py..."
+        printf "\n%s\n" "${delimiter}"
+        prepare_tcmalloc
+        accelerate launch --num_cpu_threads_per_process=6 "${LAUNCH_SCRIPT}" "$@"
+    else
+        printf "\n%s\n" "${delimiter}"
+        printf "Launching launch.py..."
+        printf "\n%s\n" "${delimiter}"
+        prepare_tcmalloc
+        "${python_cmd}" "${LAUNCH_SCRIPT}" "$@"
+    fi
+
+    if [[ ! -f tmp/restart ]]; then
+        KEEP_GOING=0
+    fi
+done