#!/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 names are read from each ONNX model at runtime rather than
# hardcoded, since they differ between models (LVFace-B: "data", arcface_r18:
# "input", arcface_w600k_{r50,mbf}: "input.1").
#
# 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 "+ $*"; "$@"; }

# The input tensor name is not the same across models — LVFace-B uses "data",
# arcface_r18 uses "input", and arcface_w600k_{r50,mbf} use "input.1". A
# hardcoded name makes trtexec fail with "Cannot find input tensor with name
# ...", so read it from the model instead.
input_name() {
    python3 - "$1" <<'PY'
import sys
try:
    import onnxruntime as ort
except ImportError:
    sys.exit("onnxruntime is required to read the model's input name")
sess = ort.InferenceSession(sys.argv[1], providers=["CPUExecutionProvider"])
print(sess.get_inputs()[0].name)
PY
}

ARCFACE_IN="$(input_name "$ARCFACE_MODEL")"
SCRFD_IN="$(input_name "$SCRFD_MODEL")"

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

echo
echo "== SCRFD == (input tensor: $SCRFD_IN)"
run trtexec \
    --onnx="$SCRFD_MODEL" \
    --fp16 \
    --minShapes="$SCRFD_IN":1x3x640x640 \
    --optShapes="$SCRFD_IN":1x3x640x640 \
    --maxShapes="$SCRFD_IN":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.
    SCENE_IN="$(input_name "$SCENE_MODEL")"
    echo "   (input tensor: $SCENE_IN)"
    run trtexec \
        --onnx="$SCENE_MODEL" \
        --fp16 \
        --minShapes="$SCENE_IN":1x100x27x48x3 \
        --optShapes="$SCENE_IN":1x100x27x48x3 \
        --maxShapes="$SCENE_IN":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."