#!/usr/bin/env bash
# ComfyUI — Stable Diffusion / SDXL runner for Apple Silicon (MPS).
#
# Why this exists: MODELBEAST's flux_local uses mflux, which is FLUX-only and
# CANNOT run SD1.5/SDXL checkpoints. ComfyUI is the engine for the ~44GB of
# Civitai models in ~/Documents/localmodels (A1111 layout).
#
# Portable across the fleet: reads models from $LOCALMODELS (default
# ~/Documents/localmodels) via extra_model_paths.yaml — no copying into vendor/.
#
#   bash scripts/install_comfyui.sh
#   cd vendor/comfyui && .venv/bin/python main.py --port 8188
set -euo pipefail
cd "$(dirname "$0")/.."
VENDOR="vendor/comfyui"
LOCALMODELS="${LOCALMODELS:-$HOME/Documents/localmodels}"

# uv lives in different places per node (brew on the Studios, ~/.local/bin on the mini)
UV="$(command -v uv || true)"
[ -x "$UV" ] || UV=/opt/homebrew/bin/uv
[ -x "$UV" ] || UV="$HOME/.local/bin/uv"
[ -x "$UV" ] || { echo "[comfy] ERROR: uv not found"; exit 1; }
echo "[comfy] uv: $UV"

[ -d "$VENDOR" ] || git clone --depth 1 https://github.com/comfyanonymous/ComfyUI.git "$VENDOR"
cd "$VENDOR"

echo "[comfy] venv (python 3.12) ..."
[ -x .venv/bin/python ] || "$UV" venv --python 3.12 .venv

echo "[comfy] torch (MPS) + requirements ..."
"$UV" pip install --python .venv/bin/python torch torchvision torchaudio
"$UV" pip install --python .venv/bin/python -r requirements.txt

echo "[comfy] pointing at models: $LOCALMODELS"
cat > extra_model_paths.yaml <<YAML
# Read the existing A1111-style model folder in place (no copying).
a111:
    base_path: $LOCALMODELS
    checkpoints: Stable-diffusion
    loras: Lora
YAML

echo "[comfy] verifying ..."
.venv/bin/python -c "import torch; print('[comfy] torch', torch.__version__, 'mps', torch.backends.mps.is_available())"
if [ -d "$LOCALMODELS/Stable-diffusion" ]; then
  echo "[comfy] checkpoints found: $(ls "$LOCALMODELS"/Stable-diffusion/*.safetensors "$LOCALMODELS"/Stable-diffusion/*.ckpt 2>/dev/null | wc -l | tr -d ' ')"
  echo "[comfy] loras found:       $(ls "$LOCALMODELS"/Lora/*.safetensors 2>/dev/null | wc -l | tr -d ' ')"
else
  echo "[comfy] NOTE: $LOCALMODELS not present yet — rsync the models, then just start ComfyUI."
fi
echo "[comfy] done — run: cd $VENDOR && .venv/bin/python main.py --port 8188"
echo "[comfy] reminder: the SD1.5 LoRAs only work with Hyper_Realism_1.2_fp16 (see localmodels/README.md)"