#!/usr/bin/env bash
# Set OMNIVAE_CODE_ROOT before running.
#
# Standard inference:
#   bash infer.sh --config-load-path configs/config_template.json --generate \
#     --prompt "A small sailboat crossing a calm lake at sunset." --seconds 5
#
# TIRx kernels on the unpruned SparseMoE checkpoint:
#   TIRX=1 bash infer.sh --config-load-path configs/config_template.json --generate \
#     --prompt "A small sailboat crossing a calm lake at sunset." --seconds 5
#
# TIRx kernels on an elastic EENP checkpoint (SwiGLU7, QK, pre-norm, gather/scatter):
#   TIRX_ELASTIC=1 bash infer.sh --config-load-path configs/config_template.json --generate \
#     --prompt "A small sailboat crossing a calm lake at sunset." --seconds 5
#
# Elastic width without TIRx kernels:
#   ELASTIC=1 bash infer.sh --config-load-path configs/config_template.json --generate \
#     --prompt "A small sailboat crossing a calm lake at sunset." --seconds 5
set -euo pipefail
REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
export PYTHONPATH="${REPO}:${OMNIVAE_CODE_ROOT:-}${PYTHONPATH:+:${PYTHONPATH}}"
export OMNIVAE_CODE_ROOT="${OMNIVAE_CODE_ROOT:-}"
export OMNIVAE_AUDIO_CKPT_DIR="${OMNIVAE_AUDIO_CKPT_DIR:-}"
export DAVINCI_DISABLE_MAGI_COMPILE="${DAVINCI_DISABLE_MAGI_COMPILE:-1}"
export MAGI_DISABLE_COMPILE="${MAGI_DISABLE_COMPILE:-1}"
export TORCHDYNAMO_DISABLE="${TORCHDYNAMO_DISABLE:-1}"
export PYTHONUNBUFFERED=1
cd "$REPO"

if [[ "${TIRX_ELASTIC:-0}" == "1" ]]; then
  export MAGI_TIRX_ELASTIC_IO="${MAGI_TIRX_ELASTIC_IO:-1}"
  export MAGI_TIRX_QK_ROPE="${MAGI_TIRX_QK_ROPE:-1}"
  export MAGI_TIRX_PRENORM="${MAGI_TIRX_PRENORM:-1}"
  export MAGI_TIRX_SWIGLU7="${MAGI_TIRX_SWIGLU7:-1}"
  export MAGI_TIRX_PREWARM="${MAGI_TIRX_PREWARM:-1}"
  exec python infer_tirx_elastic_io.py "$@"
fi

if [[ "${TIRX:-0}" == "1" ]]; then
  export MAGI_TIRX_SCATTER="${MAGI_TIRX_SCATTER:-1}"
  export MAGI_TIRX_GATHER="${MAGI_TIRX_GATHER:-1}"
  export MAGI_TIRX_DOWN="${MAGI_TIRX_DOWN:-1}"
  export MAGI_TIRX_DOWN_PREWARM="${MAGI_TIRX_DOWN_PREWARM:-1}"
  export MAGI_TIRX_QK_ROPE="${MAGI_TIRX_QK_ROPE:-1}"
  export MAGI_TIRX_PRENORM="${MAGI_TIRX_PRENORM:-1}"
  export MAGI_TIRX_SWIGLU7="${MAGI_TIRX_SWIGLU7:-1}"
  export MAGI_TIRX_PREWARM="${MAGI_TIRX_PREWARM:-1}"
  exec python infer_tirx_moe_scatter.py "$@"
fi

if [[ "${ELASTIC:-0}" == "1" ]]; then
  exec python crown/infer_entry.py "$@"
fi
exec python inference/test_infer_seedance.py "$@"