From bd2e829ae32014f613500dad3188804bbb4de6a6 Mon Sep 17 00:00:00 2001
From: JSTayco <jstayco@protonmail.ch>
Date: Fri, 31 Mar 2023 17:24:09 -0700
Subject: [PATCH] More safeties around Python ops

More safeties and more code consolidation. Now we try to exit the python venv after all the python operations. All the python operations were consolidated to facilitate this.
---
 requirements.txt |  3 ++-
 setup.sh         | 58 +++++++++++++++++++++++++++++++++---------------
 2 files changed, 42 insertions(+), 19 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 882617b..83508f5 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -13,7 +13,8 @@ lion-pytorch==0.0.6
 opencv-python==4.7.0.68
 pytorch-lightning==1.9.0
 safetensors==0.2.6
-tensorboard==2.10.1
+tensorboard==2.10.1 ; sys_platform != 'darwin'
+tensorboard==2.12.1 ; sys_platform == 'darwin'
 tk==0.1.0
 toml==0.10.2
 transformers==4.26.0
diff --git a/setup.sh b/setup.sh
index 650bf4c..14de32d 100755
--- a/setup.sh
+++ b/setup.sh
@@ -194,7 +194,22 @@ create_symlinks() {
   fi
 }
 
-install_pip_dependencies() {
+install_python_dependencies() {
+  # Switch to local virtual env
+  echo "Switching to virtual Python environment."
+  if command -v python3 >/dev/null; then
+    python3 -m venv venv
+  elif command -v python3.10 >/dev/null; then
+    python3.10 -m venv venv
+  else
+    echo "Valid python3 or python3.10 binary not found."
+    echo "Cannot proceed with the python steps."
+    return 1
+  fi
+
+  # Activate the virtual environment
+  source venv/bin/activate
+
   # Updating pip if there is one
   echo "Checking for pip updates before Python operations."
   python3 -m pip install --upgrade pip >&3
@@ -216,9 +231,23 @@ install_pip_dependencies() {
     ;;
   esac
 
+  if [ "$RUNPOD" = true ]; then
+    echo "Installing tenssort."
+    pip install tensorrt >&3
+  fi
+
   # DEBUG ONLY (Update this version number to whatever PyCharm recommends)
   # pip install pydevd-pycharm~=223.8836.43
-  python -m pip install --use-pep517 --upgrade -r requirements.txt >&3
+  python -m pip install --use-pep517 --upgrade -r "$DIR/requirements.txt" >&3
+
+  if [ -n "$VIRTUAL_ENV" ]; then
+    if command -v deactivate >/dev/null; then
+      echo "Exiting Python virtual environment."
+      deactivate
+    else
+      echo "deactivate command not found. Could still be in the Python virtual environment."
+    fi
+  fi
 }
 
 # Attempt to non-interactively install a default accelerate config file unless specified otherwise.
@@ -301,6 +330,8 @@ update_kohya_ss() {
       echo "You need to install git."
       echo "Rerun this after installing git or run this script with -n to skip the git operations."
     fi
+  else
+    echo "Skipping git operations."
   fi
 }
 
@@ -437,14 +468,10 @@ if [[ "$OSTYPE" == "linux-gnu"* ]]; then
     fi
   fi
 
-  python3 -m venv venv
-  source venv/bin/activate
-  install_pip_dependencies
+  install_python_dependencies
 
-  # We need this extra package and setup if we are running in a runpod
+  # We need just a little bit more setup for non-interactive environments
   if [ "$RUNPOD" = true ]; then
-    echo "Installing tenssort."
-    pip install tensorrt >&3
     # Symlink paths
     libnvinfer_plugin_symlink="$VENV_DIR/lib/python3.10/site-packages/tensorrt/libnvinfer_plugin.so.7"
     libnvinfer_symlink="$VENV_DIR/lib/python3.10/site-packages/tensorrt/libnvinfer.so.7"
@@ -516,17 +543,12 @@ elif [[ "$OSTYPE" == "darwin"* ]]; then
 
   update_kohya_ss
 
-  if command -v python3.10 >/dev/null; then
-    python3.10 -m venv venv
-    source venv/bin/activate
-    install_pip_dependencies
-    configure_accelerate
-    echo -e "Setup finished! Run ./gui.sh to start."
-  else
-    echo "Python not found. Please ensure you install Python."
-    echo "The brew command for Python 3.10 is: brew install python@3.10"
-    exit 1
+  if ! install_python_dependencies; then
+    echo "You may need to install Python. The command for this is brew install python@3.10."
   fi
+
+  configure_accelerate
+  echo -e "Setup finished! Run ./gui.sh to start."
 elif [[ "$OSTYPE" == "cygwin" ]]; then
   # Cygwin is a standalone suite of Linux utilies on Windows
   echo "This hasn't been validated on cygwin yet."