#!/usr/bin/env bash
# Pre-build TensorRT engines for ArcFace and SCRFD with the same shape profiles
# the runtime nodes use. First-run ORT engine builds take 30–90 s per model and
# block the pipeline; this script does it offline so cold starts are instant.
#
# Profiles must match src/arcface_embedder.hpp and src/scrfd_decoder.hpp:
#   ArcFace    : min=1x3x112x112 opt=Nx3x112x112 max=Nx3x112x112  (N = embed batch)
#   SCRFD      : 1x3x640x640 (fixed; we letterbox to this)
#   TransNetV2 : 1x100x27x48x3 (fixed; scene detector window), input tensor "input"
#
# These trtexec-built engines are *not* picked up by the ORT TRT EP cache —
# ORT uses its own engine format. The point of this script is:
#   (a) sanity-check that the ONNX models build under TRT at all;
#   (b) measure pure inference latency without ORT overhead.
# Run scene_analyze normally and ORT will populate ./trt_cache itself.
set -euo pipefail

ROOT="$(cd "$(dirname "$0")/.." && pwd)"
MODELS="$ROOT/models"
OUT="$ROOT/trt_cache"
mkdir -p "$OUT"

EMBED_BATCH="${EMBED_BATCH:-4}"
ARCFACE_MODEL="${ARCFACE_MODEL:-$MODELS/arcface_w600k_r50.onnx}"
SCRFD_MODEL="${SCRFD_MODEL:-$MODELS/scrfd_500m_bnkps.onnx}"
SCENE_MODEL="${SCENE_MODEL:-$MODELS/transnetv2.onnx}"

run() { echo "+ $*"; "$@"; }

echo "== ArcFace =="
run trtexec \
    --onnx="$ARCFACE_MODEL" \
    --fp16 \
    --minShapes=input.1:1x3x112x112 \
    --optShapes=input.1:${EMBED_BATCH}x3x112x112 \
    --maxShapes=input.1:${EMBED_BATCH}x3x112x112 \
    --saveEngine="$OUT/arcface.$(basename "$ARCFACE_MODEL" .onnx).b${EMBED_BATCH}.fp16.engine" \
    --useCudaGraph

echo
echo "== SCRFD =="
run trtexec \
    --onnx="$SCRFD_MODEL" \
    --fp16 \
    --minShapes=input.1:1x3x640x640 \
    --optShapes=input.1:1x3x640x640 \
    --maxShapes=input.1:1x3x640x640 \
    --saveEngine="$OUT/scrfd.$(basename "$SCRFD_MODEL" .onnx).640.fp16.engine" \
    --useCudaGraph

if [[ -f "$SCENE_MODEL" ]]; then
    echo
    echo "== TransNetV2 (scene detector) =="
    # Fixed 1x100x27x48x3 window. The raw-TRT scene detector backend loads this
    # engine directly via --scene-detector-engine; the ORT-TRT EP builds its own.
    run trtexec \
        --onnx="$SCENE_MODEL" \
        --fp16 \
        --minShapes=input:1x100x27x48x3 \
        --optShapes=input:1x100x27x48x3 \
        --maxShapes=input:1x100x27x48x3 \
        --saveEngine="$OUT/transnetv2.100x27x48.fp16.engine" \
        --useCudaGraph
else
    echo
    echo "== TransNetV2 skipped (no $SCENE_MODEL) =="
fi

echo
echo "Engines saved under: $OUT"
echo "Look for 'mean: ... ms' in each section for per-call latency."