From bff107878b517a220c8eb5b126fa27fdeb2d71fc Mon Sep 17 00:00:00 2001
From: JSTayco <jstayco@protonmail.ch>
Date: Wed, 29 Mar 2023 16:05:57 -0700
Subject: [PATCH] Default install location is now environment-based

Better default install locations and updated the README to reflect.
---
 README.md |  8 ++++++++
 setup.sh  | 54 ++++++++++++++++++++++++++++++++++++++----------------
 2 files changed, 46 insertions(+), 16 deletions(-)

diff --git a/README.md b/README.md
index a81de9b..dca0b1d 100644
--- a/README.md
+++ b/README.md
@@ -11,6 +11,7 @@ If you run on Linux and would like to use the GUI, there is now a port of it as
   - [Linux/macOS](#linux-and-macos-dependencies)
 - [Installation](#installation)
     - [Linux/macOS](#linux-and-macos)
+      - [Default Install Locations](#install-location)
     - [Windows](#windows)
     - [CUDNN 8.6](#optional--cudnn-86)
 - [Upgrading](#upgrading)
@@ -90,6 +91,13 @@ Options:
   -h, --help                    Show this screen.
 ```
 
+#### Install location
+
+The default install location for Linux is `/opt/kohya_ss`. If /opt is not writeable, the fallback is `$HOME/kohya_ss`. Lastly, if all else fails it will simply install to the current folder you are in.
+
+On macOS and other non-Linux machines, it will default install to `$HOME/kohya_ss` followed by where you're currently at if there's no access to $HOME.
+You can override this behavior by specifying an install directory with the -d option.
+
 If you are using the interactive mode, our default values for the accelerate config screen after running the script answer "This machine", "None", "No" for the remaining questions.
 These are the same answers as the Windows install.
 
diff --git a/setup.sh b/setup.sh
index 152eb67..112d942 100755
--- a/setup.sh
+++ b/setup.sh
@@ -24,15 +24,50 @@ Options:
 EOF
 }
 
+# Checks to see if variable is set and non-empty.
+# This is defined first, so we can use the function for some default variable values
+env_var_exists() {
+  if [[ ! -v "$1" ]] || [[ -z "$1" ]]; then
+    return 1
+  else
+    return 0
+  fi
+}
+
+# Need RUNPOD to have a default value before first access
+RUNPOD=false
+if env_var_exists RUNPOD_POD_ID || env_var_exists RUNPOD_API_KEY; then
+  RUNPOD=true
+fi
+
 # Variables defined before the getopts loop, so we have sane default values.
-DIR="/workspace/kohya_ss"
+# Default installation locations based on OS and environment
+if [[ "$OSTYPE" == "linux-gnu"* ]]; then
+  if [ "$RUNPOD" = true ]; then
+    DIR="/workspace/kohya_ss"
+  elif [ -w "/opt" ]; then
+    DIR="/opt/kohya_ss"
+  elif env_var_exists HOME; then
+    DIR="$HOME/kohya_ss"
+  else
+    # The last fallback is simply PWD
+    DIR="$(PWD)"
+  fi
+else
+  if env_var_exists HOME; then
+    DIR="$HOME/kohya_ss"
+  else
+    # The last fallback is simply PWD
+    DIR="$(PWD)"
+  fi
+fi
+
 BRANCH="master"
 GIT_REPO="https://github.com/bmaltais/kohya_ss.git"
-RUNPOD=false
 INTERACTIVE=false
 PUBLIC=false
 
-while getopts "b:d:g:ir-:" opt; do
+while getopts "b:d:g:ipr-:" opt; do
   # support long options: https://stackoverflow.com/a/28466267/519360
   if [ "$opt" = "-" ]; then # long option: reformulate OPT and OPTARG
     opt="${OPTARG%%=*}"     # extract long option name
@@ -67,15 +102,6 @@ if [[ "$OSTYPE" == "linux-gnu"* ]]; then
     root=true
   fi
 
-  # Checks to see if variable is set and non-empty.
-  env_var_exists() {
-    if [[ ! -v "$1" ]] || [[ -z "$1" ]]; then
-      return 1
-    else
-      return 0
-    fi
-  }
-
   get_distro_name() {
     local line
     if [ -f /etc/os-release ]; then
@@ -142,10 +168,6 @@ if [[ "$OSTYPE" == "linux-gnu"* ]]; then
     echo "$FREESPACEINGB"
   }
 
-  if env_var_exists RUNPOD_POD_ID || env_var_exists RUNPOD_API_KEY; then
-    RUNPOD=true
-  fi
-
   # Offer a warning and opportunity to cancel the installation if < 10Gb of Free Space detected
   if [ "$(size_available)" -lt 10 ]; then
     echo "You have less than 10Gb of free space. This installation may fail."