# 2 GPU
# accelerate launch --config_file scripts/accelerate_configs/multi_gpu.yaml --num_processes=2 --main_process_port 29501 scripts/train_omnigen.py --config config/grpo.py:general_radiomics_omnigen_full_2gpu
# KL variant:
# CONFIG_ENTRY=config/grpo.py:general_radiomics_omnigen_full_2gpu_kl bash scripts/single_node/grpo.sh
CURRENT_SERVER_IP="169.254.3.1"

if ip -o -4 addr show 2>/dev/null | awk '{split($4, address, "/"); print address[1]}' | grep -qx "$CURRENT_SERVER_IP"; then
  PROJECT_ROOT=$(CDPATH= cd -- "$(dirname -- "$0")/../.." && pwd)
  export PYTHONPATH=$PROJECT_ROOT:$PYTHONPATH
  export OMNIGEN_CODE_ROOT=/home/wenting/zr/gen_code
  export PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}
else
  export PYTHONPATH=/NAS_REMOTE/vicky/wt/codes/cxr_syn:/NAS_REMOTE/vicky/wt/codes/cxr_syn/flow_grpo:$PYTHONPATH
  export OMNIGEN_CODE_ROOT=/NAS_REMOTE/vicky/wt/codes/cxr_syn/gen_code
fi
export PYTHONPATH=$OMNIGEN_CODE_ROOT:$PYTHONPATH
export HF_TOKEN="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"
export HF_HOME="/NAS_REMOTE/vicky/wt/huggingface/models"
export HF_HUB_CACHE="/tmp/flow_grpo_hf_cache/hub"
export WANDB_API_KEY="e17ae3f7c8533a627b504b2f3a84b90d44fa3b31"

CUDA_VISIBLE_DEVICES=0,1 python3 -m accelerate.commands.launch \
  --config_file scripts/accelerate_configs/multi_gpu.yaml \
  --num_processes=${NUM_PROCESSES:-2} \
  --main_process_port ${MAIN_PROCESS_PORT:-29501} \
  scripts/train_omnigen.py \
  --config ${CONFIG_ENTRY:-config/grpo.py:general_radiomics_omnigen_full_2gpu_kl}

# ${CONFIG_ENTRY:-config/grpo.py:general_radiomics_omnigen_full_2gpu}


# 4 GPU
# accelerate launch --config_file scripts/accelerate_configs/multi_gpu.yaml --num_processes=4 --main_process_port 29501 scripts/train_omnigen.py --config config/grpo.py:general_radiomics_omnigen_full_4gpu