#!/usr/bin/env bash
set -e

DEVICE="cuda:0"
TOKENIZER_MODEL_PATH="Qwen/Qwen3-TTS-Tokenizer-12Hz"
INIT_MODEL_PATH="TTS-ORG/mans_model_mostafa/checkpoint-epoch-0"
RAW_JSONL="train_raw.jsonl"
TRAIN_JSONL="train_with_codes.jsonl"
OUTPUT_DIR="output_mostafa"
BATCH_SIZE=1
LR=2e-6
EPOCHS=4
SPEAKER_NAME="mostafa_speaker"
HF_REPO="TTS-ORG/mans_model_mostafa"

# Step 1: tokenize audio into codec codes
python prepare_data.py \
  --device ${DEVICE} \
  --tokenizer_model_path ${TOKENIZER_MODEL_PATH} \
  --input_jsonl ${RAW_JSONL} \
  --output_jsonl ${TRAIN_JSONL}

# Step 2: fine-tune with per-epoch HF push
python scripts/train_with_push.py \
  --init_model_path ${INIT_MODEL_PATH} \
  --output_model_path ${OUTPUT_DIR} \
  --train_jsonl ${TRAIN_JSONL} \
  --batch_size ${BATCH_SIZE} \
  --lr ${LR} \
  --num_epochs ${EPOCHS} \
  --speaker_name ${SPEAKER_NAME} \
  --hf_repo ${HF_REPO}