Try using TCMalloc on Linux by default

3 files changed, 22 insertions, 3 deletions

diff --git a/README.md b/README.md
index b67e2296..20f74531 100644
--- a/README.md
+++ b/README.md
@@ -120,6 +120,7 @@ sudo pacman -S wget git python3
 bash <(wget -qO- https://raw.githubusercontent.com/AUTOMATIC1111/stable-diffusion-webui/master/webui.sh)

 ```

 3. Run `webui.sh`.

+4. Check `webui-user.sh` for options.

 ### Installation on Apple Silicon

 

 Find the instructions [here](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Installation-on-Apple-Silicon).

diff --git a/webui-user.sh b/webui-user.sh
index bfa53cb7..49a426ff 100644
--- a/webui-user.sh
+++ b/webui-user.sh
@@ -43,4 +43,7 @@
 # Uncomment to enable accelerated launch
 #export ACCELERATE="True"
 
+# Uncomment to disable TCMalloc
+#export NO_TCMALLOC="True"
+
 ###########################################
diff --git a/webui.sh b/webui.sh
index 8cdad22d..1122b965 100755
--- a/webui.sh
+++ b/webui.sh
@@ -113,13 +113,13 @@ case "$gpu_info" in
         printf "Experimental support for Renoir: make sure to have at least 4GB of VRAM and 10GB of RAM or enable cpu mode: --use-cpu all --no-half"
         printf "\n%s\n" "${delimiter}"
     ;;
-    *) 
+    *)
     ;;
 esac
 if echo "$gpu_info" | grep -q "AMD" && [[ -z "${TORCH_COMMAND}" ]]
 then
     export TORCH_COMMAND="pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/rocm5.2"
-fi  
+fi
 
 for preq in "${GIT}" "${python_cmd}"
 do
@@ -172,15 +172,30 @@ else
     exit 1
 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)"
+        if [[ ! -z "${TCMALLOC}" ]]; then
+            echo "Using TCMalloc: ${TCMALLOC}"
+            export LD_PRELOAD="${TCMALLOC}"
+        else
+            printf "\e[1m\e[31mCannot locate TCMalloc (improves CPU memory usage)\e[0m\n"
+        fi
+    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}"      
+    printf "\n%s\n" "${delimiter}"
+    prepare_tcmalloc
     exec "${python_cmd}" "${LAUNCH_SCRIPT}" "$@"
 fi