#!/bin/bash

### Training parameters
# Model to train
QUALITY=high
# These two parameters will dictate, how much VRAM will be required
# NOTE: The following config will require around 20 GB VRAM (high model)
# Do not lower BATCHSIZE too much, as training will have a lot of "jitter".
# But reducing PHONEME_MAX too much will discard too long training data
# However: I was able to train on ~6-7 GB VRAM with BATCHSIZE=12 and PHONEME_MAX=400 (the jitter was clearly visible in tensorboard)
BATCHSIZE=32
PHONEME_MAX=400

# Base model to start training from. E.g. here from Thorsten Voice
BASE_CHKPOINT="./checkpoints/epoch=2665-step=1182078.ckpt"

# Start training from BASE_CHKPOINT
CHKPOINT=$BASE_CHKPOINT

# CONTINUE from an existing checkpoint.
# NOTE: It is good practice, to write down the last checkpoint, that was trained after you aborted training.
# The current checkpoint is automatically printed out by the training process.
#CHKPOINT="/training/traindata/lightning_logs/version_21/checkpoints/epoch=5661-step=1778926.ckpt"

python3 -m piper_train \
    --dataset-dir ./traindata \
    --accelerator 'gpu' \
    --gpus 1 \
    --batch-size ${BATCHSIZE} \
    --validation-split 0.0 \
    --num-test-examples 0 \
    --max_epochs 6000 \
    --resume_from_checkpoint "${CHKPOINT}" \
    --checkpoint-epochs 1 \
    --precision 16 \
    --max-phoneme-ids ${PHONEME_MAX} \
    --quality ${QUALITY}

    # Batchsize: 32
    # max-phenomene-ids: 400