#!/usr/bin/env bash
# REG/train.py：与主仓库类似，可单独指定数据根目录与预处理 cls 特征目录。
# 数据布局：${DATA_DIR}/imagenet_256_vae/ 下 VAE latent；
#          ${SEMANTIC_FEATURES_DIR}/ 下 img-feature-*.npy + dataset.json（与 parallel_encode 一致）。

NUM_GPUS=4

# ------------ 按本机路径修改 ------------
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"

# 后台示例（与主实验脚本风格一致）：
# nohup bash train.sh > jsflow-experiment.log 2>&1 &

nohup accelerate launch --multi_gpu --num_processes "${NUM_GPUS}" --mixed_precision bf16 train.py \
    --report-to wandb \
    --allow-tf32 \
    --mixed-precision bf16 \
    --seed 0 \
    --path-type linear \
    --prediction v \
    --weighting uniform \
    --model SiT-XL/2 \
    --enc-type dinov2-vit-b \
    --encoder-depth 8 \
    --proj-coeff 0.5 \
    --output-dir exps \
    --exp-name jsflow-experiment-0.75 \
    --batch-size 256 \
    --data-dir "${DATA_DIR}" \
    --semantic-features-dir "${SEMANTIC_FEATURES_DIR}" \
    --learning-rate 0.00005 \
    --t-c 0.75 \
    --cls 0.05 \
    --ot-cls \
    > jsflow-experiment.log 2>&1 & 

# 说明：
# - 不使用预处理特征、改在线抽 DINO 时：去掉 --semantic-features-dir，并保证 data-dir 为 REG 原布局
#  （imagenet_256_vae + vae-sd）。
# - 关闭 minibatch OT：追加 --no-ot-cls。
# - 主仓库 train.py 中的 --weight-ratio / --semantic-reg-coeff / --repa-* 等为本 REG 脚本未实现项；
#   投影强度请用 --proj-coeff，cls 流损失权重用 --cls。