#!/usr/bin/env bash
set -euo pipefail
sample_root=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)
PI05_CHECKPOINT_DIR=${PI05_CHECKPOINT_DIR:-$sample_root/checkpoints/pi05_fold_clothes/epoch_2}
: "${TOKENIZER_DIR:?set TOKENIZER_DIR}"
export TOKENIZER_DIR
: "${UMI_INFER_DATA:?set UMI_INFER_DATA to the calibrated cell directory}"
pi_python=${PI05_PYTHON:-python}
shadow="$sample_root/.runtime/pi05_transformers"
# Match the training launcher: patch a local shadow, never site-packages.
"$pi_python" - "$sample_root" "$shadow" <<'PY'
from pathlib import Path
import shutil,sys,transformers
root,shadow=map(Path,sys.argv[1:]);replacement=root/'openpi-umi/src/openpi/models_pytorch/transformers_replace'
if transformers.__version__!='4.53.2':raise SystemExit('The bundled OpenPI snapshot requires transformers==4.53.2')
if shadow.exists():shutil.rmtree(shadow)
shutil.copytree(Path(transformers.__file__).parent,shadow/'transformers')
shutil.copytree(replacement,shadow/'transformers',dirs_exist_ok=True)
PY
export PYTHONPATH="$shadow:$sample_root/flexiv-umi-infer:$sample_root/openpi-umi/src:$sample_root/openpi-umi/packages/openpi-client/src${PYTHONPATH:+:$PYTHONPATH}"
exec "$pi_python" -m umi_infer.policy_server --policy pi05 \
  --config pi05_umi_v3 --num-denoise-steps "${PI05_DENOISE_STEPS:-10}" --checkpoint-dir "$PI05_CHECKPOINT_DIR" \
  --eef-extrinsic-left "$UMI_INFER_DATA/E_left.npy" --eef-extrinsic-right "$UMI_INFER_DATA/E_right.npy" \
  --prompt "${TASK_PROMPT:-fold the clothes}" \
  --host "${SERVER_HOST:-127.0.0.1}" --port "${SERVER_PORT:-8765}" --device "${DEVICE:-cuda:0}" "$@"