#!/usr/bin/env bash
# Live websocket server (requires PersonaPlex weights + the moshi sources).
# Set these env vars first:
#   MOSHI_WEIGHT=/path/to/personaplex/model.safetensors
#   MIMI_WEIGHT=/path/to/mimi.safetensors
#   TOKENIZER=/path/to/tokenizer_spm_32k_3.model
#   MOSHI_SRC=/path/to/moshi_fork   (the moshi_fork/python sources)
set -e
ROOT="$(cd "$(dirname "$0")/.." && pwd)"
CODE="$ROOT/code/emo_code"
CK="$ROOT/weights_hf/imtalker/generator_v35.ckpt"
[ -f "$CK" ] || CK="$ROOT/weights_hf/imtalker/generator.ckpt"
export PYTHONPATH="${MOSHI_SRC:-}:$CODE:$CODE/generator:$CODE/renderer"
cd "$CODE"
python liveTryHeliumFrontendDequeStaticPoseFP32FM_ws_binary_AHAudioPace_EmoText.py \
  --generator_path "$CK" \
  --renderer_path "$ROOT/weights_hf/imtalker/renderer.ckpt" \
  --ref_path "$ROOT/weights_hf/weights/ref.jpeg" \
  --wav2vec_model_path "$ROOT/weights_hf/weights/wav2vec2-base-960h" \
  --audio_direct 1 --emo_enable 1 --emo_mode autotext \
  --emo_lookahead_ms 400 --head_pin 0.6 --neutral_as_happy 1 \
  --e_cfg_scale 2.5 --a_cfg_scale 4.0 --motion_crossfade_frames 5 \
  --renderer_precision bf16 --emo_inten_max 0.55 \
  --idle_still 1 --idle_rms_threshold 0.004 --idle_crossfade_frames 10 \
  --wav2vec_sec 2.0 --fm_chunk_frames 50 --audio_chunk_sec 2.0 \
  --static_pose_cam --latent_smooth 1.0 --nfe 10 --output_audio_codec pcm \
  --enable_moshi_reply --direct_reply_hidden --reply_hidden_steps_per_chunk 25 \
  --render_sub_batch 6 --prebuffer_chunks 2 --buffer_ms 1200 --drain_chunks 1 \
  --moshi_root "${MOSHI_ROOT:-$ROOT/weights_hf}" --mimi_hf_repo nvidia/personaplex-7b-v1 \
  --moshi_weight "${MOSHI_WEIGHT:?set MOSHI_WEIGHT}" \
  --mimi_weight "${MIMI_WEIGHT:?set MIMI_WEIGHT}" \
  --tokenizer "${TOKENIZER:?set TOKENIZER}" \
  --voice_prompt "${VOICE_PROMPT:-Robert_5.pt}" \
  --voice_prompt_dir "${VOICE_DIR:-$ROOT/weights_hf/weights/voices}" \
  --text_prompt "You are Robert, a warm and highly expressive voice assistant. Mirror the caller's emotion strongly." \
  --device cuda --port "${PORT:-8998}"