Create setup.sh

setup.sh

@@ -0,0 +1,182 @@
+#!/bin/bash
+set -euo pipefail
+
+# Install directly in the SSH user's home directory
+INSTALL_BASE="${HOME}"
+AI_TOOLKIT_DIR="${INSTALL_BASE}/ai-toolkit"
+UI_PORT=8675
+
+SETUP_LOG="/tmp/ai_toolkit_setup.log"
+UI_LOG="/tmp/ai_toolkit_ui.log"
+
+: > "${SETUP_LOG}"
+
+# Color formatting
+if [ -t 1 ]; then
+  RESET=$'\033[0m'
+  BOLD=$'\033[1m'
+  RED=$'\033[31m'
+  GREEN=$'\033[32m'
+  YELLOW=$'\033[33m'
+  BLUE=$'\033[34m'
+  CYAN=$'\033[36m'
+else
+  RESET=""
+  BOLD=""
+  RED=""
+  GREEN=""
+  YELLOW=""
+  BLUE=""
+  CYAN=""
+fi
+
+log() { printf "%b[%s]%b %s\n" "${BLUE}" "INFO" "${RESET}" "$1"; }
+err() { printf "%b[%s]%b %s\n" "${RED}" "ERROR" "${RESET}" "$1" >&2; exit 1; }
+ok()  { printf "%b[%s]%b %s\n" "${GREEN}" "OK" "${RESET}" "$1"; }
+section() { printf "\n%b%s%b\n" "${BOLD}${CYAN}" "$1" "${RESET}"; }
+kv()  { printf "  %b%-22s%b %s\n" "${YELLOW}" "$1" "${RESET}" "$2"; }
+
+run_logged() {
+  local label="$1"
+  local logfile="$2"
+  shift 2
+
+  printf "  %b[%s]%b %s\n" "${BLUE}" ".." "${RESET}" "$label"
+  if "$@" >>"${logfile}" 2>&1; then
+    printf "  %b[%s]%b %s\n" "${GREEN}" "OK" "${RESET}" "$label"
+  else
+    printf "  %b[%s]%b %s\n" "${RED}" "FAIL" "${RESET}" "$label" >&2
+    printf "%s\n" "--- last 30 lines from ${logfile} ---" >&2
+    tail -n 30 "${logfile}" >&2 || true
+    exit 1
+  fi
+}
+
+# ------------------ 1. WORKSPACE PREPARATION ------------------
+section "1. Preparing Workspace Directory"
+mkdir -p "${INSTALL_BASE}"
+ok "Installation home directory: ${INSTALL_BASE}"
+
+kv "Install Base (Home)" "${INSTALL_BASE}"
+kv "ai-toolkit Path"     "${AI_TOOLKIT_DIR}"
+kv "Setup Log"           "${SETUP_LOG}"
+
+# ------------------ 2. SYSTEM DEPENDENCIES & NVM NODE.JS 24 ------------------
+section "2. Installing System Dependencies & Node.js 24 via NVM"
+
+run_logged "apt update" "${SETUP_LOG}" sudo apt-get update -y || apt-get update -y
+run_logged "apt packages" "${SETUP_LOG}" sudo apt-get install -y \
+  curl git screen python3 python3-venv python3-pip \
+  zip unzip nano ffmpeg build-essential || apt-get install -y \
+  curl git screen python3 python3-venv python3-pip \
+  zip unzip nano ffmpeg build-essential
+
+# Download and install nvm & Node.js 24
+log "Installing nvm and Node.js 24..."
+export NVM_DIR="$HOME/.nvm"
+if [ ! -s "$NVM_DIR/nvm.sh" ]; then
+  run_logged "download nvm" "${SETUP_LOG}" bash -lc 'curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.6/install.sh | bash'
+fi
+
+# Load nvm into script context
+\. "$HOME/.nvm/nvm.sh"
+
+run_logged "install Node.js 24" "${SETUP_LOG}" nvm install 24
+nvm use 24
+
+log "Node.js version: $(node -v)"
+log "npm version: $(npm -v)"
+
+# Install uv for fast Python package installation
+if ! command -v uv >/dev/null 2>&1; then
+  run_logged "install uv" "${SETUP_LOG}" bash -lc 'curl -Ls https://astral.sh/uv/install.sh | sh'
+fi
+export PATH="${HOME}/.cargo/bin:/root/.cargo/bin:$PATH"
+
+ok "System dependencies & Node.js 24 installed"
+
+# ------------------ 3. AI-TOOLKIT REPO & PYTHON VENV ------------------
+section "3. Cloning ai-toolkit & Setting up Environment"
+
+cd "${INSTALL_BASE}"
+
+if [ ! -d "${AI_TOOLKIT_DIR}" ]; then
+  run_logged "git clone ai-toolkit" "${SETUP_LOG}" git clone --recursive https://github.com/ostris/ai-toolkit.git "${AI_TOOLKIT_DIR}"
+else
+  log "ai-toolkit directory exists. Fetching updates..."
+  cd "${AI_TOOLKIT_DIR}"
+  git pull || true
+fi
+
+cd "${AI_TOOLKIT_DIR}"
+run_logged "git submodules update" "${SETUP_LOG}" git submodule update --init --recursive
+
+# Create Virtual Environment
+if [ ! -d "venv" ]; then
+  run_logged "create venv" "${SETUP_LOG}" python3 -m venv venv
+fi
+
+source venv/bin/activate
+run_logged "upgrade pip" "${SETUP_LOG}" pip install --upgrade pip
+
+# Install PyTorch & Dependencies via uv using specified CUDA 12.8 wheel index
+log "Installing PyTorch 2.11.0 (CUDA 12.8)..."
+run_logged "install torch cu128" "${SETUP_LOG}" uv pip install torch==2.11.0 torchvision==0.26.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu128
+
+if [ -f "requirements.txt" ]; then
+  run_logged "install requirements.txt" "${SETUP_LOG}" uv pip install -r requirements.txt
+fi
+
+if [ -f "setup.py" ] || [ -f "pyproject.toml" ]; then
+  run_logged "install ai-toolkit package" "${SETUP_LOG}" uv pip install -e . || true
+fi
+
+run_logged "install HF hub & utilities" "${SETUP_LOG}" uv pip install \
+  huggingface_hub[cli] hf_transfer tensorboard wandb bitsandbytes diffusers transformers peft accelerate
+
+# Create standard subdirectories in ai-toolkit
+mkdir -p dataset config models output
+
+ok "Python environment setup complete"
+
+# ------------------ 4. AI-TOOLKIT WEB UI SETUP ------------------
+section "4. Setting up AI-Toolkit Web UI"
+
+if [ -d "${AI_TOOLKIT_DIR}/ui" ]; then
+  cd "${AI_TOOLKIT_DIR}/ui"
+  run_logged "npm install (ui dependencies)" "${SETUP_LOG}" npm install
+  ok "Web UI dependencies installed"
+else
+  err "ui directory not found in ai-toolkit."
+fi
+
+# ------------------ 5. LAUNCH WEB UI SCREEN SESSION ------------------
+section "5. Launching Web UI"
+
+screen -dmS aitoolkit_ui bash -lc "
+  export NVM_DIR=\"\$HOME/.nvm\"
+  [ -s \"\$NVM_DIR/nvm.sh\" ] && \. \"\$NVM_DIR/nvm.sh\"
+  nvm use 24
+  cd ${AI_TOOLKIT_DIR}/ui
+  source ../venv/bin/activate
+  export HF_HUB_ENABLE_HF_TRANSFER=1
+  echo '=== STARTING AI-TOOLKIT WEB UI (Port ${UI_PORT}) ==='
+  npm run build_and_start > ${UI_LOG} 2>&1
+  exec bash
+"
+log "Screen session 'aitoolkit_ui' launched on port ${UI_PORT}"
+
+# ------------------ SUMMARY ------------------
+section "Setup Complete!"
+kv "ai-toolkit Path" "${AI_TOOLKIT_DIR}"
+kv "Node.js Version" "$(node -v)"
+kv "PyTorch Version" "2.11.0 (cu128)"
+kv "Web UI Port"     "${UI_PORT}"
+kv "Screen Session"  "screen -r aitoolkit_ui"
+
+section "🌐 Direct Connection Info"
+echo "  • Access URL: http://<YOUR_VM_IP>:${UI_PORT}"
+echo "  • Local/SSH:  http://localhost:${UI_PORT}"
+echo "  • Attach UI:  screen -r aitoolkit_ui"
+echo ""
+ok "Done!"