update_comfy

#!/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:-<not importable>}  (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