#!/usr/bin/env bash
# Ruestet eine Blackwell-Instanz fuer Krea 2 in NVFP4 aus.
# Funktioniert auf vastai/comfy (ComfyUI vorhanden) und auf nackten PyTorch-Images.
set -euo pipefail

export HF_HOME="${HF_HOME:-/workspace/.hf}"
export HF_HUB_ENABLE_HF_TRANSFER=1
export PIP_DISABLE_PIP_VERSION_CHECK=1

echo "=== [1/5] GPU pruefen ==="
nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader
if ! nvidia-smi --query-gpu=name --format=csv,noheader | grep -qE 'RTX (50|PRO)'; then
  echo "!! WARNUNG: keine Blackwell-Karte erkannt - NVFP4 wird nicht laufen."
fi
CU=$(python -c "import torch;print(torch.version.cuda or '?')" 2>/dev/null || echo "?")
echo "torch-CUDA: $CU"
case "$CU" in 13*) : ;; *) echo "!! WARNUNG: NVFP4 braucht CUDA 13+. Gefunden: $CU (FP8 laeuft trotzdem)";; esac

echo "=== [2/5] ComfyUI finden ==="
if [ -z "${COMFY:-}" ]; then
  COMFY=$(find /workspace /opt /root -maxdepth 5 -type f -name main.py -path '*ComfyUI*' 2>/dev/null | head -1)
  COMFY=${COMFY:+$(dirname "$COMFY")}
fi
if [ -z "${COMFY:-}" ]; then
  echo "  nicht gefunden - klone nach /workspace/ComfyUI"
  COMFY=/workspace/ComfyUI
  git clone --depth 1 https://github.com/comfyanonymous/ComfyUI "$COMFY"
  pip install -q -r "$COMFY/requirements.txt"
fi
export COMFY
echo "  ComfyUI: $COMFY"

echo "=== [3/5] NVFP4-Kernel (comfy-kitchen) ==="
pip install -q --upgrade "huggingface_hub[hf_transfer]" comfy-kitchen
python -c "import comfy_kitchen as k; print('  comfy-kitchen', getattr(k,'__version__','?'))" 2>/dev/null \
  || echo "  (Import-Check uebersprungen - Paket ist installiert)"

echo "=== [4/5] Modelle laden ==="
python "$(dirname "$(readlink -f "$0")")/download_models.py"

echo "=== [5/5] Fertig ==="
echo "Modelle liegen unter $COMFY/models/"
ls -la "$COMFY/models/diffusion_models" 2>/dev/null | tail -n +2 || true
echo
if curl -sf -m 2 http://127.0.0.1:8188/system_stats >/dev/null 2>&1; then
  echo "ComfyUI laeuft bereits auf :8188 - im Browser neu laden, dann sind die Modelle da."
  echo "Messung:  python $(dirname "$(readlink -f "$0")")/bench.py --runs 5"
else
  echo "Starte ComfyUI auf :8188 ..."
  cd "$COMFY"
  exec python main.py --listen 0.0.0.0 --port 8188 --preview-method none
fi