#!/bin/bash
# One-time (idempotent) host setup for the Qwen-Image-Edit service.
# Runs as the service user (NO sudo). Safe to re-run: existing pieces are skipped.
#
# Builds, under the project BASE (the parent of this api/ dir):
#   venv/        torch 2.3.1+rocm5.7 + ComfyUI deps   (gfx906 / ROCm 5.7)
#   ComfyUI/     pinned to v0.3.77 + ComfyUI-GGUF custom node
#   ComfyUI/models/{unet,text_encoders,vae}/  the v23 Q8 GGUF + encoder + vae
set -e

source "$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )/env.sh"
GGUF_NODE="$COMFY/custom_nodes/ComfyUI-GGUF"
COMFY_TAG="v0.3.77"   # newest tag that runs on torch 2.3.1 (no comfy_kitchen)

echo "[bootstrap] BASE=$BASE  VENV=$VENV"

# --- ComfyUI (pinned) -------------------------------------------------------
if [ ! -d "$COMFY/.git" ]; then
  echo "[bootstrap] cloning ComfyUI @ $COMFY_TAG ..."
  git clone --depth 1 --branch "$COMFY_TAG" \
      https://github.com/comfyanonymous/ComfyUI.git "$COMFY"
fi

# --- venv + python deps -----------------------------------------------------
if [ ! -d "$VENV" ]; then
  echo "[bootstrap] creating venv at $VENV ..."
  python3 -m venv "$VENV"
fi
source "$VENV/bin/activate"
python -m pip install --upgrade pip wheel
echo "[bootstrap] installing torch (rocm5.7) ..."
pip install torch==2.3.1+rocm5.7 torchvision==0.18.1+rocm5.7 \
    --index-url https://download.pytorch.org/whl/rocm5.7
echo "[bootstrap] installing ComfyUI requirements ..."
pip install -r "$COMFY/requirements.txt"

# --- ComfyUI-GGUF custom node ----------------------------------------------
if [ ! -d "$GGUF_NODE" ]; then
  echo "[bootstrap] cloning ComfyUI-GGUF ..."
  git clone --depth 1 https://github.com/city96/ComfyUI-GGUF.git "$GGUF_NODE"
fi
pip install -r "$GGUF_NODE/requirements.txt" || pip install gguf

# --- API deps ---------------------------------------------------------------
pip install fastapi "uvicorn[standard]" websocket-client python-multipart pillow requests

# --- models (resume-safe; skipped if already complete) ----------------------
M="$COMFY/models"
mkdir -p "$M/unet" "$M/text_encoders" "$M/vae"
dl() {  # url  dest
  if [ -s "$2" ]; then echo "[bootstrap] have $(basename "$2")"; return; fi
  echo "[bootstrap] downloading $(basename "$2") ..."
  wget -c -q -O "$2" "$1"
}
dl "https://huggingface.co/Novice25/Qwen-Image-Edit-Rapid-AIO-GGUF/resolve/main/v23/Qwen-Rapid-NSFW-v23_Q8_0.gguf" \
   "$M/unet/Qwen-Rapid-NSFW-v23_Q8_0.gguf"
dl "https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/resolve/main/split_files/text_encoders/qwen_2.5_vl_7b_fp8_scaled.safetensors" \
   "$M/text_encoders/qwen_2.5_vl_7b_fp8_scaled.safetensors"
dl "https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/resolve/main/split_files/vae/qwen_image_vae.safetensors" \
   "$M/vae/qwen_image_vae.safetensors"

echo "[bootstrap] verifying torch + GPU ..."
python -c "import torch; print('torch', torch.__version__, 'cuda', torch.cuda.is_available(), torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'NO GPU')"
echo "[bootstrap] BOOTSTRAP_DONE"