diff options
Diffstat (limited to 'webui.sh')
| -rwxr-xr-x | webui.sh | 45 |
1 files changed, 27 insertions, 18 deletions
@@ -112,7 +112,7 @@ 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 ;; @@ -124,9 +124,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 +193,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 +203,23 @@ 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 +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 |