65 lines
2.4 KiB
Bash
65 lines
2.4 KiB
Bash
|
#!/usr/bin/env bash
|
||
|
|
set -euo pipefail
|
||
|
|
|
||
|
|
# === Config ===
|
||
|
|
INSTALL_DIR="${INSTALL_DIR:-/workspace/ai-toolkit}"
|
||
|
|
UI_PORT="${UI_PORT:-8675}" # AI-Toolkit UI default port per docs
|
||
|
|
AUTH_TOKEN="${AI_TOOLKIT_AUTH:-}" # Must be provided via environment
|
||
|
|
NODE_MAJOR="${NODE_MAJOR:-20}" # Install Node.js 20 LTS
|
||
|
|
PYTHON_BIN="${PYTHON_BIN:-python3}" # Use the system Python to create venv
|
||
|
|
|
||
|
|
if [[ -z "${AUTH_TOKEN}" ]]; then
|
||
|
|
echo "ERROR: AI_TOOLKIT_AUTH environment variable is not set."
|
||
|
|
echo "Set it in your RunPod environment (e.g., AI_TOOLKIT_AUTH=someStrongPassword) and re-run."
|
||
|
|
exit 1
|
||
|
|
fi
|
||
|
|
|
||
|
|
echo "==> Updating apt and installing dependencies..."
|
||
|
|
apt-get update -y
|
||
|
|
DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
|
||
|
|
git curl ca-certificates build-essential pkg-config \
|
||
|
|
python3-venv python3-dev \
|
||
|
|
ffmpeg libgl1 libglib2.0-0
|
||
|
|
|
||
|
|
echo "==> Installing Node.js ${NODE_MAJOR} (required: Node >=18)..."
|
||
|
|
curl -fsSL https://deb.nodesource.com/setup_${NODE_MAJOR}.x | bash -
|
||
|
|
DEBIAN_FRONTEND=noninteractive apt-get install -y nodejs
|
||
|
|
|
||
|
|
echo "==> Cloning AI-Toolkit into ${INSTALL_DIR}..."
|
||
|
|
mkdir -p "$(dirname "${INSTALL_DIR}")"
|
||
|
|
if [[ ! -d "${INSTALL_DIR}" ]]; then
|
||
|
|
git clone https://github.com/ostris/ai-toolkit.git "${INSTALL_DIR}"
|
||
|
|
else
|
||
|
|
echo " Repo already exists, pulling latest..."
|
||
|
|
git -C "${INSTALL_DIR}" pull --rebase
|
||
|
|
fi
|
||
|
|
git -C "${INSTALL_DIR}" submodule update --init --recursive
|
||
|
|
|
||
|
|
echo "==> Creating Python venv and upgrading pip..."
|
||
|
|
cd "${INSTALL_DIR}"
|
||
|
|
${PYTHON_BIN} -m venv venv
|
||
|
|
# shellcheck disable=SC1091
|
||
|
|
source venv/bin/activate
|
||
|
|
python -m pip install --upgrade pip wheel setuptools
|
||
|
|
|
||
|
|
echo "==> Installing PyTorch 2.8.0 + CUDA 12.8 wheels (torch/vision/audio)..."
|
||
|
|
# Keep versions aligned; use cu128 wheel index.
|
||
|
|
pip install --no-cache-dir \
|
||
|
|
torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 \
|
||
|
|
--index-url https://download.pytorch.org/whl/cu128
|
||
|
|
|
||
|
|
echo "==> Installing AI-Toolkit Python requirements..."
|
||
|
|
pip install --no-cache-dir -r requirements.txt
|
||
|
|
|
||
|
|
echo "==> Building and starting the AI-Toolkit UI on port ${UI_PORT}..."
|
||
|
|
# The UI script uses port 8675 by default; we enforce our env and expose host binding.
|
||
|
|
export AI_TOOLKIT_AUTH="${AUTH_TOKEN}"
|
||
|
|
export PORT="${UI_PORT}"
|
||
|
|
export HOST="0.0.0.0"
|
||
|
|
|
||
|
|
cd ui
|
||
|
|
# Use npm to install/build and then start the UI. This blocks (good for container foreground).
|
||
|
|
npm install
|
||
|
|
npm run build_and_start
|
||
|
|
|