Skip to content

Instantly share code, notes, and snippets.

@sinclairfr
Last active April 11, 2026 09:41
Show Gist options
  • Select an option

  • Save sinclairfr/6d70bcc02ad25e9ccd2c26691cfdf502 to your computer and use it in GitHub Desktop.

Select an option

Save sinclairfr/6d70bcc02ad25e9ccd2c26691cfdf502 to your computer and use it in GitHub Desktop.
#!/bin/bash
set -e
COMFYUI_DIR="/workspace/runpod-slim/ComfyUI"
VENV_DIR="$COMFYUI_DIR/.venv-cu128"
OLD_VENV_DIR="$COMFYUI_DIR/.venv"
FILEBROWSER_CONFIG="/root/.config/filebrowser/config.json"
DB_FILE="/workspace/runpod-slim/filebrowser.db"
S3_OFFLOADER_DIR="/workspace/comfyui_S3_offloader"
AI_TOOLKIT_DIR="/workspace/ai-toolkit"
# Controls whether ai-toolkit server gets launched at the end
# Set RUN_AI_TOOLKIT=true (or 1 or yes) in RunPod env vars to enable
RUN_AI_TOOLKIT="${RUN_AI_TOOLKIT:-false}"
log() {
echo "[start.sh] $*"
}
# ---------------------------------------------------------------------------- #
# Function Definitions #
# ---------------------------------------------------------------------------- #
setup_ssh() {
mkdir -p ~/.ssh
if [ ! -f /etc/ssh/ssh_host_ed25519_key ]; then
ssh-keygen -A -q
fi
if [[ -n "$PUBLIC_KEY" ]]; then
# Append only if not already present (idempotent on restarts)
grep -qxF "$PUBLIC_KEY" ~/.ssh/authorized_keys 2>/dev/null \
|| echo "$PUBLIC_KEY" >> ~/.ssh/authorized_keys
chmod 700 ~/.ssh
chmod 600 ~/.ssh/authorized_keys
log "SSH public key installed from PUBLIC_KEY env var"
else
RANDOM_PASS=$(openssl rand -base64 12)
echo "root:${RANDOM_PASS}" | chpasswd
log "Generated random SSH password for root: ${RANDOM_PASS}"
fi
grep -q "^PermitUserEnvironment yes" /etc/ssh/sshd_config \
|| echo "PermitUserEnvironment yes" >> /etc/ssh/sshd_config
/usr/sbin/sshd
}
setup_github_ssh() {
# Restores a GitHub deploy/personal key from GITHUB_SSH_KEY env var.
# Store the *private* key (base64-encoded) in RunPod env vars:
# GITHUB_SSH_KEY=$(base64 -w0 ~/.ssh/id_ed25519)
# The matching public key must be added to your GitHub account / repo.
if [[ -z "$GITHUB_SSH_KEY" ]]; then
log "GITHUB_SSH_KEY not set — skipping GitHub SSH setup"
return
fi
mkdir -p ~/.ssh
echo "$GITHUB_SSH_KEY" | base64 -d > ~/.ssh/github_key
chmod 600 ~/.ssh/github_key
# Add github.com to known_hosts to avoid interactive host confirmation
ssh-keyscan -t ed25519 github.com >> ~/.ssh/known_hosts 2>/dev/null
# Wire up the key for github.com in SSH config
if ! grep -q "Host github.com" ~/.ssh/config 2>/dev/null; then
cat >> ~/.ssh/config << 'EOF'
Host github.com
HostName github.com
User git
IdentityFile ~/.ssh/github_key
IdentitiesOnly yes
EOF
chmod 600 ~/.ssh/config
fi
log "GitHub SSH key configured (~/.ssh/github_key)"
}
export_env_vars() {
log "Exporting environment variables..."
ENV_FILE="/etc/environment"
PAM_ENV_FILE="/etc/security/pam_env.conf"
SSH_ENV_FILE="/root/.ssh/environment"
cp "$ENV_FILE" "${ENV_FILE}.bak" 2>/dev/null || true
cp "$PAM_ENV_FILE" "${PAM_ENV_FILE}.bak" 2>/dev/null || true
> "$ENV_FILE"
> "$PAM_ENV_FILE"
mkdir -p /root/.ssh
> "$SSH_ENV_FILE"
> /etc/rp_environment
printenv | grep -E '^RUNPOD_|^PATH=|^_=|^CUDA|^LD_LIBRARY_PATH|^PYTHONPATH' | while read -r line; do
name=$(echo "$line" | cut -d= -f1)
value=$(echo "$line" | cut -d= -f2-)
echo "$name=\"$value\"" >> "$ENV_FILE"
echo "$name DEFAULT=\"$value\"" >> "$PAM_ENV_FILE"
echo "$name=\"$value\"" >> "$SSH_ENV_FILE"
echo "export $name=\"$value\"" >> /etc/rp_environment
done
grep -q 'source /etc/rp_environment' ~/.bashrc 2>/dev/null \
|| echo 'source /etc/rp_environment' >> ~/.bashrc
grep -q 'source /etc/rp_environment' /etc/bash.bashrc 2>/dev/null \
|| echo 'source /etc/rp_environment' >> /etc/bash.bashrc
chmod 644 "$ENV_FILE" "$PAM_ENV_FILE"
chmod 600 "$SSH_ENV_FILE"
}
start_s3_offloader() {
log "Starting S3 offloader setup..."
if [[ ! -d "${S3_OFFLOADER_DIR}" ]]; then
log "Cloning comfyui_S3_offloader..."
git clone https://github.com/sinclairfr/comfyui_S3_offloader "${S3_OFFLOADER_DIR}" \
&& log "S3 offloader cloned" \
|| log "WARNING: S3 offloader clone failed"
else
log "comfyui_S3_offloader already present, skipping clone"
fi
if [[ -f "${S3_OFFLOADER_DIR}/app.py" ]]; then
if [[ -x "${S3_OFFLOADER_DIR}/.venv/bin/python" ]]; then
S3_PY="${S3_OFFLOADER_DIR}/.venv/bin/python"
elif [[ -x "/workspace/venv/bin/python" ]]; then
S3_PY="/workspace/venv/bin/python"
else
S3_PY="python"
fi
log "S3 offloader using: ${S3_PY}"
if ! "${S3_PY}" -c "import flask, boto3, dotenv" >/dev/null 2>&1; then
log "Installing S3 offloader deps..."
"${S3_PY}" -m pip install -r "${S3_OFFLOADER_DIR}/requirements.txt" || true
fi
nohup "${S3_PY}" "${S3_OFFLOADER_DIR}/app.py" \
> "${S3_OFFLOADER_DIR}/app.log" 2>&1 &
log "S3 offloader started (PID $!)"
else
log "S3 offloader app.py not found, skipping"
fi
}
start_ai_toolkit() {
# Resolve the Python interpreter to use — prefer ai-toolkit's own venv
if [[ -x "${AI_TOOLKIT_DIR}/.venv/bin/python" ]]; then
ATK_PY="${AI_TOOLKIT_DIR}/.venv/bin/python"
elif [[ -x "${AI_TOOLKIT_DIR}/venv/bin/python" ]]; then
ATK_PY="${AI_TOOLKIT_DIR}/venv/bin/python"
else
# Fall back to the ComfyUI venv — not ideal but better than nothing
ATK_PY="${VENV_DIR}/bin/python"
log "WARNING: no dedicated ai-toolkit venv found, falling back to ComfyUI venv"
fi
# Determine entry point — newer versions expose a proper server script
if [[ -f "${AI_TOOLKIT_DIR}/run_gradio.py" ]]; then
ATK_ENTRY="run_gradio.py"
elif [[ -f "${AI_TOOLKIT_DIR}/toolkit/ui/app.py" ]]; then
ATK_ENTRY="toolkit/ui/app.py"
else
log "WARNING: ai-toolkit entry point not found, skipping launch"
return
fi
log "Starting ai-toolkit (${ATK_ENTRY}) on port 7860..."
cd "${AI_TOOLKIT_DIR}"
nohup "${ATK_PY}" "${ATK_ENTRY}" \
> "${AI_TOOLKIT_DIR}/server.log" 2>&1 &
ATK_PID=$!
log "ai-toolkit started (PID ${ATK_PID}) — logs: ${AI_TOOLKIT_DIR}/server.log"
cd - >/dev/null
}
# ---------------------------------------------------------------------------- #
# Main Program #
# ---------------------------------------------------------------------------- #
setup_ssh
export_env_vars
setup_github_ssh
if [ ! -f "$DB_FILE" ]; then
log "Initializing FileBrowser..."
filebrowser config init
filebrowser config set --address 0.0.0.0
filebrowser config set --port 8080
filebrowser config set --root /workspace
filebrowser config set --auth.method=json
filebrowser users add admin adminadmin12 --perm.admin
else
log "Using existing FileBrowser configuration..."
fi
log "Starting FileBrowser on port 8080..."
nohup filebrowser &> /filebrowser.log &
ARGS_FILE="/workspace/runpod-slim/comfyui_args.txt"
if [ ! -f "$ARGS_FILE" ]; then
echo "# Add your custom ComfyUI arguments here (one per line)" > "$ARGS_FILE"
log "Created empty ComfyUI arguments file at $ARGS_FILE"
fi
if [ -d "$OLD_VENV_DIR" ] && [ ! -d "$VENV_DIR" ]; then
NODE_COUNT=$(find "$COMFYUI_DIR/custom_nodes" -maxdepth 2 -name "requirements.txt" 2>/dev/null | wc -l)
log "============================================="
log "CUDA 12.4 -> 12.8 migration"
log "Reinstalling deps for $NODE_COUNT custom nodes"
log "This may take several minutes"
log "============================================="
mv "$OLD_VENV_DIR" "${OLD_VENV_DIR}.bak"
cd "$COMFYUI_DIR"
python3.12 -m venv --system-site-packages "$VENV_DIR"
source "$VENV_DIR/bin/activate"
python -m ensurepip
BAKED_NODES="ComfyUI-Manager ComfyUI-KJNodes Civicomfy ComfyUI-RunpodDirect"
CURRENT=0
INSTALLED=0
for req in "$COMFYUI_DIR"/custom_nodes/*/requirements.txt; do
if [ -f "$req" ]; then
NODE_NAME=$(basename "$(dirname "$req")")
case " $BAKED_NODES " in
*" $NODE_NAME "*) continue ;;
esac
CURRENT=$((CURRENT + 1))
log "[$CURRENT] $NODE_NAME"
pip install -r "$req" 2>&1 | grep -E "^(Successfully|ERROR)" || true
INSTALLED=$((INSTALLED + 1))
fi
done
log "Upgrading ComfyUI requirements..."
pip install --upgrade -r "$COMFYUI_DIR/requirements.txt" 2>&1 | grep -E "^(Successfully|ERROR)" || true
log "Migration complete — $INSTALLED user nodes processed (${NODE_COUNT} total, baked nodes skipped)"
log "Old venv backed up at ${OLD_VENV_DIR}.bak — delete it to free space:"
log "rm -rf ${OLD_VENV_DIR}.bak"
fi
if [ ! -d "$COMFYUI_DIR" ] || [ ! -d "$VENV_DIR" ]; then
log "First time setup: Copying baked ComfyUI to workspace..."
if [ ! -d "$COMFYUI_DIR" ]; then
cp -r /opt/comfyui-baked "$COMFYUI_DIR"
log "ComfyUI copied to workspace"
fi
if [ ! -d "$VENV_DIR" ]; then
cd "$COMFYUI_DIR"
python3.12 -m venv --system-site-packages "$VENV_DIR"
source "$VENV_DIR/bin/activate"
python -m ensurepip
log "Base packages (torch, numpy, etc.) available from system site-packages"
log "ComfyUI ready — all dependencies pre-installed in image"
fi
else
source "$VENV_DIR/bin/activate"
log "Using existing ComfyUI installation"
fi
python -m pip --version > /dev/null 2>&1
start_s3_offloader
cd "$COMFYUI_DIR"
FIXED_ARGS="--listen 0.0.0.0 --port 8188"
if [ -s "$ARGS_FILE" ]; then
CUSTOM_ARGS=$(grep -v '^#' "$ARGS_FILE" | tr '\n' ' ')
if [ -n "$CUSTOM_ARGS" ]; then
FIXED_ARGS="$FIXED_ARGS $CUSTOM_ARGS"
fi
fi
log "Starting ComfyUI with args: $FIXED_ARGS"
python main.py $FIXED_ARGS &
COMFY_PID=$!
# Launch ai-toolkit if requested
case "${RUN_AI_TOOLKIT,,}" in
true|1|yes)
start_ai_toolkit
;;
*)
log "ai-toolkit disabled (RUN_AI_TOOLKIT=${RUN_AI_TOOLKIT})"
;;
esac
trap "kill $COMFY_PID 2>/dev/null" SIGTERM SIGINT
wait $COMFY_PID || true
log "============================================="
log "ComfyUI crashed — check the logs above."
log "SSH and FileBrowser are still available."
log "To restart after fixing:"
log "cd $COMFYUI_DIR && source .venv-cu128/bin/activate"
log "python main.py $FIXED_ARGS"
log "============================================="
sleep infinity
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment