akhaliq/Dramabox

1#!/bin/bash2# Launch DramaBox IC-LoRA training. Wraps src/train.py with accelerate.3# Usage:4#   ./scripts/train.sh --config configs/training_args.yaml --gpus 2,3,4,5,65set -e6 7CONFIG=""8GPUS=${GPUS:-0,1,2,3,4,5,6,7}9NUM_PROCS=${NUM_PROCS:-}10TRAIN_VAL_GPU=${TRAIN_VAL_GPU:-}11EXTRA_ARGS=""12 13while [[ $# -gt 0 ]]; do14  case $1 in15    --config)        CONFIG="$2"; shift 2;;16    --gpus)          GPUS="$2";   shift 2;;17    --num-procs)     NUM_PROCS="$2"; shift 2;;18    --train-val-gpu) TRAIN_VAL_GPU="$2"; shift 2;;19    *)               EXTRA_ARGS+=" $1"; shift;;20  esac21done22 23if [[ -z "$CONFIG" ]]; then24  echo "Usage: $0 --config <yaml> [--gpus 2,3,4,5,6] [--num-procs N] [--train-val-gpu N]"25  exit 126fi27 28REPO="$(cd "$(dirname "$0")/.." && pwd)"29PYTHON="${PYTHON:-/usr/bin/env python}"30 31# Default num_procs to gpu count if unset.32if [[ -z "$NUM_PROCS" ]]; then33  NUM_PROCS=$(awk -F',' '{print NF}' <<< "$GPUS")34fi35 36# Convert YAML config -> CLI args (accepts a flat dict mapping arg-name -> value37# and a `data_dir` / `speaker_index` list).38get() { "$PYTHON" -c "import yaml,sys; c=yaml.safe_load(open('$CONFIG')); v=c.get('$1', '$2'); print(v if not isinstance(v,(list,tuple)) else ' '.join(map(str,v)))"; }39 40DATA_DIRS=$(get data_dir "")41SPK_IDX=$(get speaker_index "")42OUT_DIR=$(get output_dir "tts_iclora_v1")43[[ "$OUT_DIR" != /* ]] && OUT_DIR="$REPO/$OUT_DIR"44CKPT=$(get checkpoint "$REPO/ltx-2.3-22b-dev.safetensors")45FULL_CKPT=$(get full_checkpoint "$REPO/ltx-2.3-22b-dev.safetensors")46BASE_MODEL=$(get base_model dev)47LORA_RANK=$(get lora_rank 128)48LORA_ALPHA=$(get lora_alpha 128)49LORA_DROPOUT=$(get lora_dropout 0.0)50RESUME_LORA=$(get resume_lora "")51[[ -n "$RESUME_LORA" && "$RESUME_LORA" != /* ]] && RESUME_LORA="$REPO/$RESUME_LORA"52REF_RATIO=$(get ref_ratio 0.3)53MAX_REF=$(get max_ref_tokens 200)54TEXT_DROP=$(get text_dropout 0.4)55STEPS=$(get steps 10000)56LR=$(get lr 0.0001)57SCHED=$(get lr_scheduler cosine)58BATCH=$(get batch_size 1)59GRAD_ACC=$(get grad_accum 4)60GRAD_NORM=$(get max_grad_norm 1.0)61SAVE_EVERY=$(get save_every 500)62LOG_EVERY=$(get log_every 50)63SEED=$(get seed 53)64WARMUP=$(get warmup_steps 500)65VAL_CFG=$(get val_config "")66[[ -n "$VAL_CFG" && "$VAL_CFG" != /* ]] && VAL_CFG="$REPO/configs/$VAL_CFG"67 68mkdir -p "$OUT_DIR"69 70CMD=( "$PYTHON" -u -m accelerate.commands.launch71      --num_processes="$NUM_PROCS" --mixed_precision=bf1672      "$REPO/src/train.py"73      --data-dir $DATA_DIRS74      --speaker-index $SPK_IDX75      --output-dir "$OUT_DIR"76      --checkpoint "$CKPT" --full-checkpoint "$FULL_CKPT" --base-model "$BASE_MODEL"77      --lora-rank "$LORA_RANK" --lora-alpha "$LORA_ALPHA" --lora-dropout "$LORA_DROPOUT"78      --ref-ratio "$REF_RATIO" --max-ref-tokens "$MAX_REF" --text-dropout "$TEXT_DROP"79      --steps "$STEPS" --lr "$LR" --lr-scheduler "$SCHED"80      --batch-size "$BATCH" --grad-accum "$GRAD_ACC" --max-grad-norm "$GRAD_NORM"81      --save-every "$SAVE_EVERY" --log-every "$LOG_EVERY" --seed "$SEED"82      --warmup-steps "$WARMUP" )83[[ -n "$RESUME_LORA" ]] && CMD+=( --resume-lora "$RESUME_LORA" )84[[ -n "$VAL_CFG"     ]] && CMD+=( --val-config "$VAL_CFG" )85CMD+=( $EXTRA_ARGS )86 87LAUNCH_ENV=( "CUDA_VISIBLE_DEVICES=$GPUS" )88[[ -n "$TRAIN_VAL_GPU" ]] && LAUNCH_ENV+=( "TRAIN_VAL_GPU=$TRAIN_VAL_GPU" )89 90echo "==== launching DramaBox training ===="91echo "  GPUs:  $GPUS  (procs: $NUM_PROCS)"92echo "  out:   $OUT_DIR"93echo "  ckpt:  $CKPT"94echo "  steps: $STEPS  lr: $LR  $SCHED  warmup: $WARMUP"95echo "======================================"96env "${LAUNCH_ENV[@]}" "${CMD[@]}"97