Upload provision.sh

provision.sh

@@ -0,0 +1,241 @@
+#!/usr/bin/env bash
+# ---------------------------------------------------------------------------
+# ComfyUI provisioning for RunPod  --  Flux.2 Klein FaceSwap + Z-Image LoRA
+#
+# RunPod template -> "Container Start Command":
+#   bash -c "curl -sL <RAW_URL_OF_THIS_FILE> -o /provision.sh && bash /provision.sh 2>&1 | tee /workspace/provision.log"
+#
+# Required template env var:
+#   HF_TOKEN          read-only HuggingFace token (Klein 9b fp8 is a gated repo)
+# Optional:
+#   START_JUPYTER=1   also run JupyterLab on :8888   (default 0)
+#   SKIP_MODELS=1     set up nodes only, no model downloads
+#   COMFY_PORT=8188   ComfyUI port
+# ---------------------------------------------------------------------------
+set -uo pipefail
+
+ROOT="${WORKSPACE:-/workspace}"
+COMFY="$ROOT/ComfyUI"
+STAGE="$ROOT/.dl"
+CONS="$ROOT/.torch-constraints.txt"
+PORT="${COMFY_PORT:-8188}"
+M="$COMFY/models"
+
+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' "$*"; }
+die()  { printf '\n\033[1;31m[fail] %s\033[0m\n' "$*"; exit 1; }
+
+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
+
+# ------------------------------------------------------------ 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: $HF_BIN"
+
+# -------------------------------------------------------------- ComfyUI ----
+if [ ! -d "$COMFY/.git" ]; then
+    log "Cloning ComfyUI"
+    git clone --depth 1 https://github.com/comfyanonymous/ComfyUI "$COMFY" || die "clone failed"
+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
+            warn "clone failed: $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
+    warn "TORCH CHANGED -> $TORCH_NOW  (expected $(head -1 "$CONS"))"
+fi
+
+# --------------------------------------------------------------- models ----
+mkdir -p "$M/diffusion_models" "$M/text_encoders" "$M/vae" "$M/loras" "$M/sams" "$M/ultralytics/bbox"
+
+fetch() {   # repo  path-in-repo  dest  expected-bytes
+    local repo="$1" path="$2" dest="$3" want="$4" have
+    mkdir -p "$(dirname "$dest")"
+    if [ -f "$dest" ]; then
+        have=$(stat -c %s "$dest")
+        if [ "$have" = "$want" ]; then
+            ok "cached  $(basename "$dest")"
+            return 0
+        fi
+        warn "size mismatch ($have != $want) -> refetch $(basename "$dest")"
+        rm -f "$dest"
+    fi
+    printf '    downloading %s ...\n' "$(basename "$dest")"
+    rm -rf "$STAGE"
+    if ! "$HF_BIN" download "$repo" "$path" --local-dir "$STAGE" >/dev/null 2>&1; then
+        warn "DOWNLOAD FAILED  $repo :: $path"
+        rm -rf "$STAGE"
+        return 1
+    fi
+    mv -f "$STAGE/$path" "$dest"
+    rm -rf "$STAGE"
+    have=$(stat -c %s "$dest")
+    if [ "$have" = "$want" ]; then
+        ok "ok      $(basename "$dest")  ($have B)"
+    else
+        warn "BAD SIZE $(basename "$dest"): $have != $want"
+        return 1
+    fi
+}
+
+fetch_url() {   # url  dest  expected-bytes
+    local url="$1" dest="$2" want="$3" have
+    mkdir -p "$(dirname "$dest")"
+    if [ -f "$dest" ]; then
+        have=$(stat -c %s "$dest")
+        if [ "$have" = "$want" ]; then
+            ok "cached  $(basename "$dest")"
+            return 0
+        fi
+        rm -f "$dest"
+    fi
+    printf '    downloading %s ...\n' "$(basename "$dest")"
+    aria2c -x8 -s8 -q --allow-overwrite=true \
+           -d "$(dirname "$dest")" -o "$(basename "$dest")" "$url" >/dev/null 2>&1 \
+        || curl -sL "$url" -o "$dest" \
+        || { warn "DOWNLOAD FAILED  $url"; return 1; }
+    have=$(stat -c %s "$dest")
+    if [ "$have" = "$want" ]; then
+        ok "ok      $(basename "$dest")  ($have B)"
+    else
+        warn "BAD SIZE $(basename "$dest"): $have != $want"
+    fi
+}
+
+if [ "${SKIP_MODELS:-0}" = "1" ]; then
+    log "Models: skipped (SKIP_MODELS=1)"
+else
+    log "Models  (~37.6 GB total)"
+
+    # -- Flux.2 Klein 9b stage (GATED: needs HF_TOKEN + accepted licence) -----
+    if [ -z "${HF_TOKEN:-}" ]; then
+        warn "HF_TOKEN is empty - the gated Klein 9b fp8 file will fail."
+        warn "Set HF_TOKEN in the template AND accept the licence once at"
+        warn "https://huggingface.co/black-forest-labs/FLUX.2-klein-9b-fp8"
+    fi
+    fetch black-forest-labs/FLUX.2-klein-9b-fp8 \
+          flux-2-klein-9b-fp8.safetensors \
+          "$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
+
+    # -- Z-Image Turbo stage (FaceDetailer) ----------------------------------
+    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
+
+    # -- your character LoRA -------------------------------------------------
+    fetch cocount21/Robert01 \
+          Robert01_epoch50.safetensors \
+          "$M/loras/Robert01_epoch50.safetensors" 174083128
+
+    # -- detectors for FaceDetailer ------------------------------------------
+    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
+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
+
+# --------------------------------------------------------------- launch ----
+printf '\n\033[1;32m==================================================\033[0m\n'
+printf '\033[1;32m  provisioning finished in %s s\033[0m\n' "$(( $(date +%s) - START_TS ))"
+printf '\033[1;32m  ComfyUI -> port %s\033[0m\n' "$PORT"
+printf '\033[1;32m==================================================\033[0m\n\n'
+
+cd "$COMFY"
+exec python main.py --listen 0.0.0.0 --port "$PORT"