#!/bin/bash # update_comfy — fix the "generation hangs at 0% / Model Initializing" problem. # # Root cause (see ./start_comfy startup log): # torch is built for CUDA 12.8 (2.x.x+cu128) but the optimized VRAM stack # (comfy-aimdo + comfy-kitchen, "DynamicVRAM" / async weight offloading) # wants CUDA 13.0 ("You need pytorch with cu130 or higher..."). aimdo hooks # cudaMallocAsync and then deadlocks the moment the sampler starts its # forward pass, so models load but 0/8 never advances. # # Two fixes, matching what you picked ("both — try disable first"): # ./update_comfy Step 1: bypass the cu130-only path by launching # with DynamicVRAM + async offload disabled. Your # H100 (80GB VRAM / 2TB RAM) does not need them — # the 38GB model fits with room to spare. No # downloads, instantly reversible. # ./update_comfy --cu130 Step 2 (optional): reinstall PyTorch built for # CUDA 13.0 so the optimized path actually works, # then re-enable it. Downloads several GB. # ./update_comfy --revert Undo: restore the stock launch (optimizer on). # ./update_comfy --status Just print the health report, change nothing. # # NOTE: /workspace/start_comfy is re-pulled from Hugging Face on every pod # boot, which wipes the Step-1 patch. Re-run ./update_comfy after each boot, # or set WORKSPACE_REFRESH=false in the pod env to stop the boot refresh. set -uo pipefail WS="/workspace" START="$WS/start_comfy" FLAGS="--disable-dynamic-vram --disable-async-offload" CU130_INDEX="https://download.pytorch.org/whl/cu130" # ---- colors (fall back to plain if not a tty) -------------------------------- if [[ -t 1 ]]; then B=$'\e[1m'; G=$'\e[32m'; Y=$'\e[33m'; R=$'\e[31m'; C=$'\e[36m'; X=$'\e[0m' else B=; G=; Y=; R=; C=; X=; fi say() { printf '%s\n' "$*"; } ok() { printf '%s\n' " ${G}✔${X} $*"; } warn() { printf '%s\n' " ${Y}!${X} $*"; } err() { printf '%s\n' " ${R}x${X} $*" >&2; } # ---- helpers ----------------------------------------------------------------- torch_cuda() { python -c "import torch;print(torch.version.cuda or 'none')" 2>/dev/null; } torch_ver() { python -c "import torch;print(torch.__version__)" 2>/dev/null; } health_report() { say "" say "${B}== ComfyUI health ==${X}" local tv tc drv tv=$(torch_ver); tc=$(torch_cuda) drv=$(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -1) say " torch : ${tv:-} (CUDA ${tc:-?})" say " driver : ${drv:-?}" say " comfy-aimdo : $(pip show comfy-aimdo 2>/dev/null | awk '/^Version/{print $2}')" say " comfy-kitchen : $(pip show comfy-kitchen 2>/dev/null | awk '/^Version/{print $2}')" if [[ "$tc" == "none" || -z "$tc" ]]; then warn "torch reports no CUDA build." elif [[ "${tc%%.*}" -ge 13 ]]; then ok "torch CUDA (${tc}) satisfies the optimized VRAM stack (needs 13.0+)." else warn "torch CUDA (${tc}) < 13.0 — optimized VRAM stack (aimdo) will hang." fi if grep -q -- "--disable-dynamic-vram" "$START" 2>/dev/null; then ok "start_comfy is patched: optimizer disabled (safe launch)." else say " start_comfy : stock (optimizer enabled)" fi say "" } apply_flags() { [[ -f "$START" ]] || { err "$START not found."; return 1; } if grep -q -- "--disable-dynamic-vram" "$START"; then ok "start_comfy already patched — nothing to do." return 0 fi cp "$START" "$START.bak.$(date +%s)" 2>/dev/null || true # append the disable flags to the main.py launch line if grep -qE 'python[0-9]* +main\.py' "$START"; then sed -i -E "/python[0-9]* +main\.py/ s/\$/ ${FLAGS}/" "$START" chmod 755 "$START" ok "Patched start_comfy → launches with: ${C}${FLAGS}${X}" else err "Could not find the 'python main.py' line in $START; not modified." return 1 fi } revert_flags() { [[ -f "$START" ]] || { err "$START not found."; return 1; } if grep -q -- "--disable-dynamic-vram" "$START"; then sed -i -E "s/ --disable-dynamic-vram --disable-async-offload//g" "$START" ok "Reverted start_comfy to stock launch (optimizer enabled)." else ok "start_comfy already stock — nothing to revert." fi } upgrade_cu130() { say "${B}== Step 2: reinstall PyTorch for CUDA 13.0 ==${X}" warn "This downloads several GB and may bump the torch version. Ctrl-C to abort." say " index: $CU130_INDEX" say "" if ! pip install --upgrade torch torchvision torchaudio --index-url "$CU130_INDEX"; then err "pip install failed — leaving Step-1 disable-flags in place so ComfyUI still works." return 1 fi local tc; tc=$(torch_cuda) if [[ -z "$tc" || "${tc%%.*}" -lt 13 ]]; then err "After install, torch CUDA is '${tc:-none}' (< 13). Keeping disable-flags." return 1 fi # sanity: torch can actually see the GPU if ! python -c "import torch;assert torch.cuda.is_available();torch.zeros(8,device='cuda').sum().item()" 2>/dev/null; then err "torch cu130 installed but a basic CUDA op failed. Keeping disable-flags for safety." return 1 fi ok "torch now on CUDA ${tc} and a CUDA op succeeded." revert_flags # optimizer can be re-enabled now that cu130 is present ok "Optimized VRAM stack re-enabled. Restart ComfyUI to use it." } # ---- main -------------------------------------------------------------------- say "${B}update_comfy${X} — ComfyUI hang fixer" case "${1:-}" in --status) health_report exit 0 ;; --revert) revert_flags health_report exit 0 ;; --cu130) # ensure we're in a working state first, then attempt the upgrade apply_flags upgrade_cu130 health_report say "Start ComfyUI with: ${C}cd $WS && ./start_comfy${X}" exit 0 ;; ""|--fix) apply_flags health_report say "Now start ComfyUI: ${C}cd $WS && ./start_comfy${X}" say "If it still hangs, escalate: add ${C}--disable-cuda-malloc${X} to the launch," say "or run ${C}./update_comfy --cu130${X} for the full CUDA-13 fix." say "" warn "start_comfy is restored from Hugging Face on every pod boot —" warn "re-run ./update_comfy after a reboot (or set WORKSPACE_REFRESH=false)." exit 0 ;; -h|--help) sed -n '2,34p' "$0"; exit 0 ;; *) err "Unknown option: $1"; sed -n '2,34p' "$0"; exit 1 ;; esac