#!/bin/bash
# apply_patches.sh — apply local fixes to the pinned ComfyUI core that upstream
# hasn't merged yet. Idempotent; safe to run on every boot / every start.
#
# Patch 1 — ComfyUI #14573: sampler hangs at "Model Initializing" (and slow
#   loads) with comfy-aimdo / DynamicVRAM on NVIDIA. Root cause:
#   model_management.py forces MIN_WEIGHT_MEMORY_RATIO = 0.0 on NVIDIA, and
#   aimdo's custom-allocator VRAM is untracked by get_free_memory(), so the
#   weight-memory formula allocates ~0 MB to the model → deadlock on the first
#   forward pass. Fix: use 0.4 (the non-NVIDIA default) so DynamicVRAM gets real
#   VRAM → no hang AND fast staged loads. Remove once upstream fixes #14573:
#   https://github.com/Comfy-Org/ComfyUI/issues/14573

set -uo pipefail
COMFYUI_DIR="${COMFYUI_DIR:-/workspace/ComfyUI}"
MM="$COMFYUI_DIR/comfy/model_management.py"

log(){ printf '[apply_patches] %s\n' "$*"; }

patch_min_weight_ratio(){
    [ -f "$MM" ] || { log "model_management.py not found ($MM) — skipping"; return 0; }
    if ! grep -qE '^[[:space:]]*MIN_WEIGHT_MEMORY_RATIO = 0\.0[[:space:]]*$' "$MM"; then
        if grep -q '#14573' "$MM"; then
            log "patch 1 already applied — ok"
        else
            log "patch 1: '= 0.0' line not present (ComfyUI changed?) — skipping, verify upstream"
        fi
        return 0
    fi
    cp "$MM" "$MM.orig.$(date +%s)" 2>/dev/null || true
    sed -i -E \
      's/^([[:space:]]*)MIN_WEIGHT_MEMORY_RATIO = 0\.0[[:space:]]*$/\1MIN_WEIGHT_MEMORY_RATIO = 0.4  # patched: ComfyUI #14573 (aimdo\/DynamicVRAM hang)/' \
      "$MM"
    if grep -q '#14573' "$MM" && ! grep -qE '^[[:space:]]*MIN_WEIGHT_MEMORY_RATIO = 0\.0[[:space:]]*$' "$MM"; then
        log "patch 1 applied: MIN_WEIGHT_MEMORY_RATIO 0.0 → 0.4"
    else
        log "WARNING: patch 1 did not apply cleanly — inspect $MM"
    fi
}

patch_min_weight_ratio
log "done"