#!/usr/bin/env bash # =========================================================================== # ComfyUI provisioning for RunPod # Flux.2 Klein head-swap + Z-Image Turbo face restore # =========================================================================== # # RunPod template -> "Container Start Command": # bash -c "curl -sL https://huggingface.co/cocount21/comfy-faceswap-setup/resolve/main/provision.sh -o /provision.sh && bash /provision.sh 2>&1 | tee /workspace/provision.log" # # ------------------------------- SETTINGS ------------------------------- # HF_TOKEN your OWN HuggingFace read token. # Needed only when WANT_FLUX=1 (Klein 9b is a gated repo). # Leave it EMPTY in the shared template - every user # types in their own at deploy time. # # WANT_FLUX=1 Flux.2 Klein 9b stage (~19 GB, gated, non-commercial) # WANT_ZIMAGE=1 Z-Image Turbo stage (~21 GB, Apache-2.0, open) # WANT_DETECTORS=1 face detector + SAM (~0.4 GB, needed by FaceDetailer) # # EXTRA_LORAS="" your own LoRAs, comma separated, "repo:path-in-repo" # e.g. EXTRA_LORAS="myname/myrepo:my_lora.safetensors" # # GET_WORKFLOWS=1 auto-install the .json workflows into ComfyUI # START_JUPYTER=0 also run JupyterLab on :8888 # SKIP_MODELS=0 set up nodes only, download nothing # COMFY_PORT=8188 # COMFY_REF="" pin ComfyUI to a git commit/tag instead of latest # =========================================================================== set -uo pipefail ROOT="${WORKSPACE:-/workspace}" COMFY="$ROOT/ComfyUI" STAGE="$ROOT/.dl" CONS="$ROOT/.torch-constraints.txt" M="$COMFY/models" PORT="${COMFY_PORT:-8188}" WANT_FLUX="${WANT_FLUX:-1}" WANT_ZIMAGE="${WANT_ZIMAGE:-1}" WANT_DETECTORS="${WANT_DETECTORS:-1}" EXTRA_LORAS="${EXTRA_LORAS:-}" GET_WORKFLOWS="${GET_WORKFLOWS:-1}" SETUP_REPO="${SETUP_REPO:-cocount21/comfy-faceswap-setup}" GATED_REPO="black-forest-labs/FLUX.2-klein-9b-fp8" GATED_FILE="flux-2-klein-9b-fp8.safetensors" REPORT=() PROBLEMS=0 log() { printf '\n\033[1;36m==> %s\033[0m\n' "$*"; } ok() { printf ' \033[1;32m%s\033[0m\n' "$*"; } warn() { printf ' \033[1;33m[warn] %s\033[0m\n' "$*"; } err() { printf ' \033[1;31m[ERROR] %s\033[0m\n' "$*"; PROBLEMS=$((PROBLEMS+1)); } die() { printf '\n\033[1;31m[fail] %s\033[0m\n' "$*"; exit 1; } banner() { printf '\n\033[1;31m' printf '#%.0s' {1..74}; printf '\n' while IFS= read -r l; do printf '# %-69s#\n' "$l"; done <<< "$1" printf '#%.0s' {1..74}; printf '\033[0m\n\n' } mkdir -p "$ROOT"; cd "$ROOT" START_TS=$(date +%s) # --------------------------------------------------------------- ssh ------- if [ -n "${PUBLIC_KEY:-}" ]; then mkdir -p ~/.ssh && chmod 700 ~/.ssh echo "$PUBLIC_KEY" >> ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys ssh-keygen -A >/dev/null 2>&1 /usr/sbin/sshd >/dev/null 2>&1 && ok "sshd up" fi # ========================================================================== # PREFLIGHT - check the gated repo BEFORE downloading tens of gigabytes # ========================================================================== if [ "${SKIP_MODELS:-0}" != "1" ] && [ "$WANT_FLUX" = "1" ]; then log "Preflight: access to the gated Flux repo" CODE=$(curl -s -o /dev/null -w '%{http_code}' -L -r 0-0 \ -H "Authorization: Bearer ${HF_TOKEN:-nope}" \ "https://huggingface.co/$GATED_REPO/resolve/main/$GATED_FILE" 2>/dev/null) case "$CODE" in 200|206) ok "access confirmed (HTTP $CODE)" ;; 401) banner "NO ACCESS TO THE FLUX MODEL (HTTP 401 - bad or missing token) Fix, takes 2 minutes: 1. Create a HuggingFace account: https://huggingface.co/join 2. Make a READ token: https://huggingface.co/settings/tokens Create new token -> type Read -> copy it 3. Put it into this pod's HF_TOKEN environment variable and restart. Skipping the Flux stage for now - Z-Image will still be installed." WANT_FLUX=0; PROBLEMS=$((PROBLEMS+1)) ;; 403) banner "FLUX LICENCE NOT ACCEPTED (HTTP 403) Your token works, but you have not accepted the model licence yet. 1. Open (logged in): https://huggingface.co/black-forest-labs/FLUX.2-klein-9b-fp8 2. Press the Agree button on the grey banner. 3. Restart this pod. Note: FLUX.2 Klein is under a NON-COMMERCIAL licence. Skipping the Flux stage for now - Z-Image will still be installed." WANT_FLUX=0; PROBLEMS=$((PROBLEMS+1)) ;; *) warn "unexpected HTTP $CODE - will try anyway" ;; esac fi # ------------------------------------------------------------ apt deps ----- if [ ! -f "$ROOT/.apt-done" ]; then log "System packages" export DEBIAN_FRONTEND=noninteractive apt-get update -qq >/dev/null 2>&1 apt-get install -y -qq git aria2 libgl1 libglib2.0-0t64 >/dev/null 2>&1 \ || apt-get install -y -qq git aria2 libgl1 libglib2.0-0 >/dev/null 2>&1 \ || warn "apt partially failed (opencv-based nodes may complain)" touch "$ROOT/.apt-done"; ok "done" fi # ------------------------------------------------- pin torch, get hf cli --- log "Python tooling" python - > "$CONS" <<'PYEOF' import torch print("torch==" + torch.__version__) for m in ("torchvision", "torchaudio"): try: print(m + "==" + __import__(m).__version__) except Exception: pass PYEOF ok "torch pinned: $(head -1 "$CONS")" pip install -q -U pip >/dev/null 2>&1 pip install -q -U "huggingface_hub[hf_transfer,cli]" >/dev/null 2>&1 export HF_HUB_ENABLE_HF_TRANSFER=1 HF_BIN="$(command -v hf || command -v huggingface-cli || true)" [ -n "$HF_BIN" ] || die "huggingface cli not found" ok "hf cli ready" # -------------------------------------------------------------- ComfyUI --- if [ ! -d "$COMFY/.git" ]; then log "Cloning ComfyUI" if [ -n "${COMFY_REF:-}" ]; then git clone https://github.com/comfyanonymous/ComfyUI "$COMFY" >/dev/null 2>&1 \ && git -C "$COMFY" checkout "$COMFY_REF" >/dev/null 2>&1 \ && ok "pinned to $COMFY_REF" || die "clone/checkout failed" else git clone --depth 1 https://github.com/comfyanonymous/ComfyUI "$COMFY" >/dev/null 2>&1 \ || die "clone failed" fi fi log "ComfyUI requirements" pip install -q -r "$COMFY/requirements.txt" -c "$CONS" >/dev/null 2>&1 || warn "some reqs failed" ok "done" # --------------------------------------------------------- custom nodes --- NODES=( "https://github.com/ltdrdata/ComfyUI-Manager" "https://github.com/ltdrdata/ComfyUI-Impact-Pack" "https://github.com/ltdrdata/ComfyUI-Impact-Subpack" "https://github.com/yolain/ComfyUI-Easy-Use" "https://github.com/rgthree/rgthree-comfy" "https://github.com/zhangp365/ComfyUI-utils-nodes" ) log "Custom nodes" mkdir -p "$COMFY/custom_nodes" for url in "${NODES[@]}"; do name="$(basename "$url")" dir="$COMFY/custom_nodes/$name" if [ ! -d "$dir/.git" ]; then if git clone --depth 1 "$url" "$dir" >/dev/null 2>&1; then ok "cloned $name" else err "clone failed: $name"; REPORT+=("NODE MISSING $name"); continue fi else ok "present $name" fi if [ -f "$dir/requirements.txt" ]; then pip install -q -r "$dir/requirements.txt" -c "$CONS" >/dev/null 2>&1 || warn "pip failed: $name" fi if [ -f "$dir/install.py" ]; then ( cd "$dir" && python install.py >/dev/null 2>&1 ) || warn "install.py failed: $name" fi done TORCH_NOW="$(python -c 'import torch; print(torch.__version__)' 2>/dev/null || echo BROKEN)" if [ "torch==$TORCH_NOW" = "$(head -1 "$CONS")" ]; then ok "torch intact ($TORCH_NOW)" else err "TORCH CHANGED -> $TORCH_NOW (expected $(head -1 "$CONS"))" REPORT+=("TORCH was replaced by a custom node - generation may fail") fi # ============================================================== models ===== mkdir -p "$M/diffusion_models" "$M/text_encoders" "$M/vae" "$M/loras" "$M/sams" "$M/ultralytics/bbox" hf_size() { # repo path -> expected bytes (0 if unknown) python - "$1" "$2" <<'PYEOF' import sys, os, json, urllib.request repo, path = sys.argv[1], sys.argv[2] try: req = urllib.request.Request( "https://huggingface.co/api/models/%s/tree/main?recursive=true" % repo) tok = os.environ.get("HF_TOKEN") if tok: req.add_header("Authorization", "Bearer " + tok) for e in json.load(urllib.request.urlopen(req, timeout=30)): if e.get("path") == path and e.get("type") == "file": print((e.get("lfs") or {}).get("size") or e.get("size") or 0) break else: print(0) except Exception: print(0) PYEOF } fetch() { # repo path-in-repo dest expected-bytes|auto local repo="$1" path="$2" dest="$3" want="$4" have base base="$(basename "$dest")" [ "$want" = "auto" ] && want="$(hf_size "$repo" "$path")" mkdir -p "$(dirname "$dest")" if [ -f "$dest" ]; then have=$(stat -c %s "$dest") if [ "$want" = "0" ] || [ "$have" = "$want" ]; then ok "cached $base"; REPORT+=("OK $base"); return 0 fi warn "size mismatch ($have != $want) -> refetch $base" rm -f "$dest" fi printf ' downloading %s ...\n' "$base" rm -rf "$STAGE" if ! "$HF_BIN" download "$repo" "$path" --local-dir "$STAGE" >/dev/null 2>&1; then err "download failed: $repo :: $path"; REPORT+=("MISSING $base") rm -rf "$STAGE"; return 1 fi mv -f "$STAGE/$path" "$dest"; rm -rf "$STAGE" have=$(stat -c %s "$dest") if [ "$want" = "0" ] || [ "$have" = "$want" ]; then ok "ok $base ($have B)"; REPORT+=("OK $base") else err "BAD SIZE $base: $have != $want"; REPORT+=("CORRUPT $base") return 1 fi } fetch_url() { # url dest expected-bytes local url="$1" dest="$2" want="$3" have base base="$(basename "$dest")" mkdir -p "$(dirname "$dest")" if [ -f "$dest" ] && [ "$(stat -c %s "$dest")" = "$want" ]; then ok "cached $base"; REPORT+=("OK $base"); return 0 fi rm -f "$dest" printf ' downloading %s ...\n' "$base" aria2c -x8 -s8 -q --allow-overwrite=true \ -d "$(dirname "$dest")" -o "$base" "$url" >/dev/null 2>&1 \ || curl -sL "$url" -o "$dest" \ || { err "download failed: $url"; REPORT+=("MISSING $base"); return 1; } have=$(stat -c %s "$dest") if [ "$have" = "$want" ]; then ok "ok $base ($have B)"; REPORT+=("OK $base") else err "BAD SIZE $base: $have != $want"; REPORT+=("CORRUPT $base") fi } if [ "${SKIP_MODELS:-0}" = "1" ]; then log "Models: skipped (SKIP_MODELS=1)" else if [ "$WANT_FLUX" = "1" ]; then log "Models: Flux.2 Klein 9b (~19 GB)" fetch "$GATED_REPO" "$GATED_FILE" \ "$M/diffusion_models/Flux-2-Klein-9b-fp8.safetensors" 9433061528 fetch Comfy-Org/vae-text-encorder-for-flux-klein-9b \ split_files/text_encoders/qwen_3_8b_fp8mixed.safetensors \ "$M/text_encoders/qwen_3_8b_fp8mixed.safetensors" 8664848742 fetch Comfy-Org/vae-text-encorder-for-flux-klein-9b \ split_files/vae/flux2-vae.safetensors \ "$M/vae/flux2-vae.safetensors" 336211292 fetch bond12321/bfs-head-klein9b \ bfs_head_v1_flux-klein_9b_step3500_rank128.safetensors \ "$M/loras/bfs_head_v1_flux-klein_9b_step3500_rank128.safetensors" 662729912 else log "Models: Flux stage skipped" REPORT+=("SKIPPED Flux.2 Klein stage") fi if [ "$WANT_ZIMAGE" = "1" ]; then log "Models: Z-Image Turbo (~21 GB)" fetch Comfy-Org/z_image_turbo \ split_files/diffusion_models/z_image_turbo_bf16.safetensors \ "$M/diffusion_models/z_image_turbo_bf16.safetensors" 12309866400 fetch Comfy-Org/z_image_turbo \ split_files/text_encoders/qwen_3_4b.safetensors \ "$M/text_encoders/qwen_3_4b.safetensors" 8044982048 fetch Comfy-Org/z_image_turbo \ split_files/vae/ae.safetensors \ "$M/vae/ae.safetensors" 335304388 else log "Models: Z-Image stage skipped" REPORT+=("SKIPPED Z-Image stage") fi if [ "$WANT_DETECTORS" = "1" ]; then log "Models: face detectors" fetch Bingsu/adetailer face_yolov8m.pt \ "$M/ultralytics/bbox/face_yolov8m.pt" 52026019 fetch_url https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth \ "$M/sams/sam_vit_b_01ec64.pth" 375042383 # Impact-Subpack refuses to load .pt detectors that are not whitelisted. # Without this FaceDetailer fails at generation time, not at startup. WL="$COMFY/user/default/ComfyUI-Impact-Subpack/model-whitelist.txt" mkdir -p "$(dirname "$WL")" touch "$WL" grep -qxF "face_yolov8m.pt" "$WL" 2>/dev/null || echo "face_yolov8m.pt" >> "$WL" ok "whitelisted face_yolov8m.pt for Impact-Subpack" fi # -- user supplied LoRAs ------------------------------------------------- if [ -n "$EXTRA_LORAS" ]; then log "Models: your own LoRAs" IFS=',' read -ra ITEMS <<< "$EXTRA_LORAS" for item in "${ITEMS[@]}"; do item="$(echo "$item" | xargs)" [ -z "$item" ] && continue lrepo="${item%%:*}"; lpath="${item#*:}" if [ "$lrepo" = "$item" ] || [ -z "$lpath" ]; then err "bad EXTRA_LORAS entry: '$item' (expected repo:file.safetensors)" continue fi fetch "$lrepo" "$lpath" "$M/loras/$(basename "$lpath")" auto done fi fi # ----------------------------------------------------------- workflows ---- if [ "$GET_WORKFLOWS" = "1" ]; then log "Workflows" WFDIR="$COMFY/user/default/workflows" mkdir -p "$WFDIR" rm -rf "$STAGE" if "$HF_BIN" download "$SETUP_REPO" --include "workflows/*.json" \ --local-dir "$STAGE" >/dev/null 2>&1 && [ -d "$STAGE/workflows" ]; then cp -f "$STAGE"/workflows/*.json "$WFDIR"/ 2>/dev/null \ && ok "installed: $(ls -1 "$WFDIR" | tr '\n' ' ')" \ || warn "no .json found in $SETUP_REPO/workflows/" else warn "no workflows/ folder in $SETUP_REPO - load the .json by hand" fi rm -rf "$STAGE" fi # ------------------------------------------------------------- jupyter ---- if [ "${START_JUPYTER:-0}" = "1" ] && command -v jupyter >/dev/null 2>&1; then JT="${JUPYTER_PASSWORD:-$(python -c 'import secrets; print(secrets.token_hex(16))')}" nohup jupyter lab --allow-root --no-browser --ip=0.0.0.0 --port=8888 \ --ServerApp.token="$JT" --ServerApp.allow_origin='*' \ --notebook-dir="$ROOT" >"$ROOT/jupyter.log" 2>&1 & log "JupyterLab :8888 token = $JT" fi # ======================================================== final report ==== ELAPSED=$(( $(date +%s) - START_TS )) printf '\n\n' printf '\033[1m=========================== SUMMARY ===========================\033[0m\n' for line in "${REPORT[@]}"; do case "$line" in OK*) printf ' \033[1;32m%s\033[0m\n' "$line" ;; SKIPPED*) printf ' \033[1;34m%s\033[0m\n' "$line" ;; *) printf ' \033[1;31m%s\033[0m\n' "$line" ;; esac done printf '\033[1m=================================================================\033[0m\n' printf ' time: %s s\n' "$ELAPSED" if [ "$PROBLEMS" -eq 0 ]; then printf '\n\033[1;32m ALL GOOD - ComfyUI is starting on port %s\033[0m\n' "$PORT" printf '\033[1;32m Connect -> HTTP Service -> :%s\033[0m\n\n' "$PORT" else printf '\n\033[1;33m %s PROBLEM(S) ABOVE - ComfyUI still starts, but some\n' "$PROBLEMS" printf ' nodes in the workflow will be red / models missing.\n' printf ' Full log: %s/provision.log\033[0m\n\n' "$ROOT" fi cd "$COMFY" exec python main.py --listen 0.0.0.0 --port "$PORT"