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

cd /home/nvidia/SiT-Complementary

export WANDB_KEY
WANDB_KEY="$(python -c 'import netrc; print(netrc.netrc().authenticators("api.wandb.ai")[2])')"

# Match the base and rotation-layer runs while keeping a separate W&B run.
export WANDB_MODE=offline
export WANDB_DIR=/data/nvidia/SiT-conv-layer-bs256/wandb
export SIT_FID_COMPARISON_OUTPUT_DIR=/home/nvidia/SiT-comparisons/bs256-lr1e-4-800ep
mkdir -p "$WANDB_DIR"

exec torchrun \
  --nnodes=1 \
  --nproc_per_node=8 \
  train_conv.py \
  --model SiT-S/2 \
  --epochs 800 \
  --data-path /home/nvidia/datasets/imagenet-1k/train \
  --results-dir /data/nvidia/SiT-conv-layer-bs256/results-800ep \
  --global-batch-size 256 \
  --learning-rate 0.0001 \
  --global-seed 0 \
  --vae ema \
  --num-workers 4 \
  --log-every 100 \
  --ckpt-every 50000 \
  --sample-every 10000 \
  --cfg-scale 4.0 \
  --run-name SiT-S-2-ConvLayer-bs256-lr1e-4-800ep \
  --fid-every-checkpoint \
  --fid-every 250000 \
  --fid-num-samples 50000 \
  --fid-reference /home/nvidia/evaluation/reference/discon-download/VIRTUAL_imagenet256_labeled.npz \
  --fid-history /data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/fid_cfg1_50k.tsv \
  --fid-per-proc-batch-size 64 \
  --fid-inception-batch-size 128 \
  --fid-num-workers 8 \
  --fid-sampling-steps 250 \
  --fid-seed 0 \
  --fid-stop-consecutive-increases 3 \
  --fid-stop-min-absolute-rise 0.25 \
  --fid-stop-min-relative-rise 0.005 \
  --wandb