#!/bin/bash

set -e

CKPT_DIR="/piper/lightning_logs"
VOICE_NAME=${VOICE_NAME:-pt_BR-well}

latest_checkpoint() {
  ls -t ${CKPT_DIR}/version_*/checkpoints/*.ckpt | head -n1
}

if compgen -G "${CKPT_DIR}/version_*/checkpoints/*.ckpt" > /dev/null; then
  CKPT=$(latest_checkpoint)
  RESUME_ARG="--ckpt_path $CKPT"

  echo "Resuming: $CKPT"
fi

python3 -m piper.train fit \
  --data.voice_name $VOICE_NAME \
  --data.espeak_voice pt-br \
  --data.audio_dir /data/wav/ \
  --data.batch_size 16 \
  --data.cache_dir /data/.cache/ \
  --data.config_path /data/${VOICE_NAME}-medium.onnx.json \
  --data.csv_path /data/metadata.csv \
  --model.sample_rate 22050 \
  --trainer.check_val_every_n_epoch 1 \
  --trainer.max_epochs 10000 \
  $RESUME_ARG

CKPT=$(latest_checkpoint)

python3 -m piper.train.export_onnx \
  --checkpoint $CKPT \
  --output-file /data/${VOICE_NAME}-medium.onnx