#!/bin/bash
# One-time (idempotent) host setup for the Qwen-Image-Edit service on the CUDA box.
# Runs as the service user (NO sudo). Safe to re-run: existing pieces are skipped.
#
# Builds, under $BASE (default ~/comfyui, outside the git repo):
#   venv/        CUDA torch + ComfyUI deps   (RTX A6000 / sm_86)
#   ComfyUI/     latest master + ComfyUI-GGUF custom node
#   ComfyUI/models/{unet,text_encoders,vae}/  the v23 Q8 GGUF + encoder + vae
#
# CUDA lifts every MI50 workaround: no rocm wheel, no v0.3.77 pin (that pin only
# existed because rocm-torch 2.3.1 lacks torch.library.custom_op), no lowvram.
set -e

source "$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )/env.sh"
GGUF_NODE="$COMFY/custom_nodes/ComfyUI-GGUF"

echo "[bootstrap] BASE=$BASE  VENV=$VENV"
mkdir -p "$BASE"

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

# --- venv + python deps -----------------------------------------------------
# Build from the real system interpreter, NOT whatever venv is currently active.
if [ ! -d "$VENV" ]; then
  echo "[bootstrap] creating venv at $VENV ..."
  /usr/bin/python3.13 -m venv "$VENV"
fi
source "$VENV/bin/activate"
python -m pip install --upgrade pip wheel
echo "[bootstrap] installing torch (CUDA cu124) ..."
# torchaudio MUST come from the cu124 index too. Latest ComfyUI imports torchaudio
# at startup (comfy/sd.py -> audio_vae); the PyPI default build links libcudart.so.13
# (CUDA 13) and crashes against our cu124 torch. The matching +cu124 build fixes it.
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
echo "[bootstrap] installing ComfyUI requirements ..."
pip install -r "$COMFY/requirements.txt"
# ComfyUI's requirements can re-pull a mismatched torchaudio from PyPI; re-pin it.
pip install torchaudio --index-url https://download.pytorch.org/whl/cu124

# --- 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) ----------------------
# Identical files/URLs to tour-comfy/bootstrap.sh.
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 --show-progress -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"
