Last active
April 11, 2026 09:41
-
-
Save sinclairfr/6d70bcc02ad25e9ccd2c26691cfdf502 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #!/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