#!/bin/bash
# =============================================================
# ComfyUI Launch Script for AMD BC-250 (ROCm / gfx1013)
# Trusted source: BC250 ROCm Install README.md
# =============================================================

set -euo pipefail

echo "=========================================="
echo " ComfyUI — BC-250 ROCm Launcher"
echo "=========================================="

# --- GPU Health Check ---
if dmesg 2>/dev/null | tail -50 | grep -qi "KIQ fence timeout"; then
    echo "[ABORT] KIQ fence timeout detected in dmesg — reboot required!"
    exit 1
fi
echo "[OK] GPU health check passed"

# --- ROCm Environment for BC-250 (gfx1013 mapped to gfx1010) ---
# Per BC250 ROCm Install docs — ONLY these vars
export HSA_OVERRIDE_GFX_VERSION=10.1.0
export HSA_ENABLE_SDMA=0
export HIP_VISIBLE_DEVICES=0
export ROCM_PATH=/opt/rocm
export HSA_TOOLS_LIB=""
export HSA_TOOLS_REPORT_LOAD_FAILURE=0
export PATH="/opt/rocm/bin:$PATH"
export LD_LIBRARY_PATH="/opt/rocm/lib"

# --- Unset harmful old workarounds ---
unset GPU_MAX_HW_QUEUES 2>/dev/null || true
unset HIP_LAUNCH_BLOCKING 2>/dev/null || true
unset GGML_CUDA_ENABLE_UNIFIED_MEMORY 2>/dev/null || true
unset GGML_HIP_HOST_ALLOC 2>/dev/null || true
unset GGML_CUDA_NO_PINNED 2>/dev/null || true
unset GGML_HIP_NO_COARSE_GRAIN 2>/dev/null || true
unset HSA_DISABLE_FRAGMENT_ALLOCATOR 2>/dev/null || true

# --- PyTorch ROCm tuning ---
export PYTORCH_HIP_ALLOC_CONF="expandable_segments:False"

# --- CPU thread optimization for VAE and other CPU-bound ops ---
# BC-250 has 12 threads — use them all
export OMP_NUM_THREADS=12
export MKL_NUM_THREADS=12
export NUMEXPR_NUM_THREADS=12
export OPENBLAS_NUM_THREADS=12

echo "[OK] ROCm environment configured"
echo "  GFX Override: $HSA_OVERRIDE_GFX_VERSION"
echo "  OMP Threads: $OMP_NUM_THREADS"

# --- Activate venv ---
cd ~/ComfyUI
source venv/bin/activate

# --- Launch ComfyUI ---
echo "[START] Launching ComfyUI on http://0.0.0.0:8188"
echo "=========================================="
python main.py \
    --listen 0.0.0.0 \
    --port 8188 \
    --force-fp32 \
    --lowvram \
    "$@"