#!/bin/bash
# VIPER training launcher -- one stage per invocation.
#   scripts/train.sh <stage> [extra args...]
#
# Uses torchrun over the 8 GPUs of one node.  Ranks alternate between the
# low-noise and high-noise Wan2.2 experts (see viper/train.py) so both halves of
# the timestep range are covered without holding 56GB of DiT on one card.
set -e
cd "$(dirname "$0")/.."
source env.sh

STAGE=$1; shift
NGPU=${NGPU:-8}

case $STAGE in
  1) EXTRA="--stage 1 --steps ${STEPS:-3000}" ;;
  2) EXTRA="--stage 2 --steps ${STEPS:-6000} --init_from checkpoints/stage1/final.pt" ;;
  3) EXTRA="--stage 3 --steps ${STEPS:-6000} --init_from checkpoints/stage2/final.pt" ;;
  *) echo "usage: $0 {1|2|3}"; exit 1 ;;
esac

exec torchrun --nproc_per_node=$NGPU --master_port=${MASTER_PORT:-29571} \
  -m viper.train $EXTRA "$@"