#!/usr/bin/env bash
# New/REG/train.sh：只对齐数据加载相关目录/参数；不引入 t_c/ot-cls。

set -euo pipefail

NUM_GPUS=4

# ------------------------- 按需修改这些路径 -------------------------
# New/REG/dataset.py 期望 data_dir 下至少有：
#   - imagenet_256_vae/
#   - vae-sd/（非预处理 semantic 模式需要）
# 若启用预处理语义：还需要
#   - imagenet_256_features/dinov2-vit-b_tmp/gpu0/dataset.json + 对应 .npy
DATA_DIR="/gemini/space/zhaozy/dataset/Imagenet/imagenet_256"
SEMANTIC_FEATURES_DIR="/gemini/space/zhaozy/dataset/Imagenet/imagenet_256/imagenet_256_features/dinov2-vit-b_tmp/gpu0"

OUTPUT_BASE_DIR="/your_path/reg_xlarge_dinov2_base_align_8_cls"
EXP_NAME="linear-dinov2-b-enc8"

# ------------------------- 不建议改动以下逻辑 -------------------------
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
TRAIN_PY="${SCRIPT_DIR}/train.py"

nohup accelerate launch --multi_gpu --num_processes "${NUM_GPUS}" "${TRAIN_PY}" \
  --report-to="wandb" \
  --allow-tf32 \
  --mixed-precision="fp16" \
  --seed=0 \
  --path-type="linear" \
  --prediction="v" \
  --weighting="uniform" \
  --model="SiT-XL/2" \
  --enc-type="dinov2-vit-b" \
  --proj-coeff=0.5 \
  --encoder-depth=8 \
  --cls=0.03 \
  --output-dir="${OUTPUT_BASE_DIR}" \
  --exp-name="${EXP_NAME}" \
  --batch-size=256 \
  --data-dir="${DATA_DIR}" \
  --semantic-features-dir="${SEMANTIC_FEATURES_DIR}" \
   > jsflow-experiment.log 2>&1 &