Blame for provision.sh

#!/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"