#!/usr/bin/env bash
# Launch a ROMA real-time demo on the 4x Quadro RTX 6000. Run INSIDE the container (workdir /app).
#
# Usage:
#   bash scripts/rtx6000/run_demo.sh [proactive|narration|mme]
#
# Demos (default: proactive):
#   proactive  - real-time proactive event alert (Speak Head fires above threshold 0.6)
#   narration  - real-time streaming narration (threshold 0.975)
#   mme        - reactive multimodal QA
#
# Turing-safe defaults (override via env before calling):
#   ROMA_DTYPE=float16   (set bfloat16 if fp16 produces NaN gate probs; slower on Turing)
#   ROMA_ATTN=sdpa       (set eager if the transformers fork rejects sdpa)
#   ROMA_LOAD_8BIT=0     (fp16 sharded across all visible GPUs; see note below)
#
# GPUs: by default ALL visible GPUs are used (device_map="auto" shards the ~22GB fp16 model across
# the 4 cards). To restrict to fewer GPUs, export CUDA_VISIBLE_DEVICES and enable 8-bit, e.g.:
#   CUDA_VISIBLE_DEVICES=0 ROMA_LOAD_8BIT=1 bash scripts/rtx6000/run_demo.sh proactive
#
# Then open http://<host>:${GRADIO_SERVER_PORT:-7860}
set -euo pipefail

# Run from the repo root so the demos' "from src.llamafactory ..." imports resolve.
cd "$(dirname "$0")/../.."

DEMO="${1:-proactive}"
export ROMA_DTYPE="${ROMA_DTYPE:-float16}"
export ROMA_ATTN="${ROMA_ATTN:-sdpa}"
export ROMA_LOAD_8BIT="${ROMA_LOAD_8BIT:-0}"
export GRADIO_SERVER_NAME="${GRADIO_SERVER_NAME:-0.0.0.0}"
export GRADIO_SERVER_PORT="${GRADIO_SERVER_PORT:-7860}"

case "${DEMO}" in
  proactive) SCRIPT="gradio/proactive_gradio.py" ;;
  narration) SCRIPT="gradio/narration_gradio.py" ;;
  mme)       SCRIPT="gradio/mme_gradio.py" ;;
  *) echo "Unknown demo '${DEMO}'. Use: proactive | narration | mme" >&2; exit 1 ;;
esac

MODEL_PATH="${ROMA_MODEL_PATH:-whole_model/model}"
if [[ ! -e "${MODEL_PATH}/config.json" ]]; then
  echo "ERROR: no model at '${MODEL_PATH}'. Run scripts/rtx6000/download_model.sh first." >&2
  exit 1
fi

echo "Launching '${DEMO}' demo (dtype=${ROMA_DTYPE} attn=${ROMA_ATTN} 8bit=${ROMA_LOAD_8BIT}) -> http://<host>:${GRADIO_SERVER_PORT}"
exec python "${SCRIPT}"