#!/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 run's low-overhead W&B recording setup. This is a new run
# because its name (and therefore deterministic W&B run ID) is unique.
export WANDB_MODE=offline
export WANDB_DIR=/home/nvidia/SiT-rot-layer-bs256/wandb
mkdir -p "$WANDB_DIR"

exec torchrun \
  --nnodes=1 \
  --nproc_per_node=8 \
  train_rot_layer.py \
  --model SiT-S/2 \
  --epochs 800 \
  --data-path /home/nvidia/datasets/imagenet-1k/train \
  --results-dir /home/nvidia/SiT-rot-layer-bs256/results-200ep \
  --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-RotLayer-bs256-lr1e-4-200ep \
  --ckpt /home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/1000800.pt \
  --fid-every-checkpoint \
  --fid-every 250000 \
  --fid-num-samples 50000 \
  --fid-reference /home/nvidia/evaluation/reference/discon-download/VIRTUAL_imagenet256_labeled.npz \
  --fid-history /home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/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