#!/usr/bin/env bash
set -euo pipefail

DATA_ROOT="${DATA_ROOT:-./datasets/celeba-hq-256}"
SAVE_DIR="${SAVE_DIR:-./runs/celebahq256_rtm}"
NUM_GPUS="${NUM_GPUS:-4}"
DINO_CACHE_DIR="${DINO_CACHE_DIR:-./dinov2_cache}"

export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"

mkdir -p "$SAVE_DIR" "$DINO_CACHE_DIR"

CKPT_DIR="$SAVE_DIR/train"
RESTORE_ARGS=()
if [ -f "$CKPT_DIR/latest-model.th" ]; then
    echo "[train_celebahq256] resuming from $CKPT_DIR/latest-model.th"
    RESTORE_ARGS+=(--restore_path "$CKPT_DIR/latest-model.th")
    [ -f "$CKPT_DIR/latest-model-ema.th" ] && RESTORE_ARGS+=(--restore_ema_path "$CKPT_DIR/latest-model-ema.th")
    [ -f "$CKPT_DIR/latest-opt.th" ]       && RESTORE_ARGS+=(--restore_optimizer_path "$CKPT_DIR/latest-opt.th")
    [ -f "$CKPT_DIR/latest-sched.th" ]     && RESTORE_ARGS+=(--restore_scheduler_path "$CKPT_DIR/latest-sched.th")
    [ -f "$CKPT_DIR/latest-log.jsonl" ]    && RESTORE_ARGS+=(--restore_log_path "$CKPT_DIR/latest-log.jsonl")
fi

torchrun --nnodes=1 --nproc_per_node="$NUM_GPUS" --standalone train.py \
    --hps celebahq256 \
    --save_dir "$SAVE_DIR" \
    --data_root "$DATA_ROOT" \
    --num_epochs 6000 \
    --use_se True \
    --residual_type convex \
    --latent_dim 512 \
    --dec_blocks '1x1,4m1,4x2,8m4,8x4,16m8,16x5,32m16,32x5,64m32,64x5,128m64,128x4,256m128,256x1' \
    --custom_width_str '256:64,128:128,64:512,32:768,16:768,8:768,4:768,1:768' \
    --force_factor 5 \
    --imle_force_resample 2 \
    --lr 0.0003 \
    --search_type lpips \
    --n_batch 24 \
    --imle_batch 128 \
    --iters_per_save 1000 \
    --iters_per_images 5000 \
    --iters_per_ckpt 500000 \
    --use_rtm True \
    --H_cycles 4 --L_cycles 2 --refinement_steps 4 \
    --num_tokens 4 \
    --rtm_hidden_size 256 \
    --dino_cache_dir "$DINO_CACHE_DIR" \
    "${RESTORE_ARGS[@]}"