#!/bin/bash
# ──────────────────────────────────────────────────────────────────────────────
# Training script for aesthetic scoring with custom score tokens.
#
# Each image is expanded into 14 independent samples:
#   - 13 per-criterion samples  (score range 1–10, bins=crit_score_bins)
#   - 1  total-score sample     (score range 0–100, bins=total_score_bins)
#
# Assistant output format (per sample):
#   <think>{reason}</think>
#   The score is <scoreK>          ← criteria sample
#
#   The score is <totalK>          ← total sample
#
# Loss = CE_loss + crit_reg_weight * Huber(criteria) + total_reg_weight * Huber(total)
# ──────────────────────────────────────────────────────────────────────────────
# export CUDA_VISIBLE_DEVICES=0
# ── Distributed configuration (auto-detect) ──────────────────────────────────
MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
MASTER_PORT=${MASTER_PORT:-$(shuf -i 20001-29999 -n 1)}
NNODES=${WORLD_SIZE:-1}

if [ -z "$NPROC_PER_NODE" ]; then
    if command -v nvidia-smi &> /dev/null; then
        NPROC_PER_NODE=$(nvidia-smi --list-gpus | wc -l)
    else
        NPROC_PER_NODE=1
    fi
fi


export DS_BUILD_SHM_COMM=0  # Disable building shared memory communication ops
export DS_BUILD_OPS=0  # Disable building ops to avoid compilation errors
# ── Paths ─────────────────────────────────────────────────────────────────────
# Base model: Qwen3-VL-8B or a previously fine-tuned checkpoint
# LLM=/mnt/hdfs/zw04mlnn01/checkpoint/llm_platform/hldy/model/Qwen/Qwen3-VL-8B-Instruct/main
LLM=
# Image root directory (must contain all image_path entries in track_1_train.json)
IMAGE_DIR=zijielou/PortraitCraft/main

# Annotation file (image_path uses subdir format: images_00/xxx.jpg)
ANNOTATION=PortraitCraft/track_1_train_subdir.json

# Output
RUN_NAME="qwen3vl_aesthetic_score_tokens_0509"
OUTPUT_DIR=result/portrait/${RUN_NAME}

# DeepSpeed config
DEEPSPEED=./scripts/zero2.json

# ── Hyper-parameters ──────────────────────────────────────────────────────────
LR=2e-5
BATCH_SIZE=2               # per GPU
GRAD_ACCUM=16                # effective batch = NPROC_PER_NODE * BATCH_SIZE * GRAD_ACCUM
NUM_EPOCHS=3

# ── Score token settings ──────────────────────────────────────────────────────
CRIT_BINS=20                # <score1>…<score20> covering 1.0–10.0
TOTAL_BINS=101              # <total1>…<total101> covering 0–100

# ── Loss weights ──────────────────────────────────────────────────────────────
CE_WEIGHT=1.0
CRIT_REG_WEIGHT=0.5         # regression loss weight for criteria samples
TOTAL_REG_WEIGHT=0.5        # regression loss weight for total-score samples
REG_LOSS_TYPE=huber         # huber | mse

# ── Image resolution ──────────────────────────────────────────────────────────
MAX_PIXELS=$((32 * 32 * 1024))   # ~451k pixels (~672px side)
MIN_PIXELS=$((32 * 32 * 224))    # ~12k  pixels
IMAGE_MAX_TOKEN_NUM=1024
# ── Entry point ───────────────────────────────────────────────────────────────
ENTRY=qwenvl/train/train_qwen.py

# ── Update aesthetic data path in data registry ───────────────────────────────
# Patch the image_dir at runtime via environment variable so we don't hard-code
# it inside __init__.py.  The AestheticScoringDataset reads IMAGE_DIR if set.
export AESTHETIC_IMAGE_DIR=${IMAGE_DIR}
export AESTHETIC_ANNOTATION=${ANNOTATION}
export WANDB_DISABLED=true

# ── Auto-compute save_steps = half epoch ──────────────────────────────────────
NUM_IMAGES=$(python -c "import json; print(len(json.load(open('${ANNOTATION}'))))")
SAMPLES_PER_IMAGE=14   # 13 criteria + 1 total
TOTAL_SAMPLES=$((NUM_IMAGES * SAMPLES_PER_IMAGE))
EFFECTIVE_BATCH=$((NPROC_PER_NODE * NNODES * BATCH_SIZE * GRAD_ACCUM))
STEPS_PER_EPOCH=$(( (TOTAL_SAMPLES + EFFECTIVE_BATCH - 1) / EFFECTIVE_BATCH ))
SAVE_STEPS=$(( STEPS_PER_EPOCH / 2 ))
[ ${SAVE_STEPS} -lt 1 ] && SAVE_STEPS=1
echo "[auto] ${NUM_IMAGES} images × ${SAMPLES_PER_IMAGE} = ${TOTAL_SAMPLES} samples"
echo "[auto] effective_batch=${EFFECTIVE_BATCH}, steps/epoch=${STEPS_PER_EPOCH}, save_steps=${SAVE_STEPS} (half epoch)"

# ── Training arguments ────────────────────────────────────────────────────────
ARGS="
    --model_name_or_path ${LLM} \
    --use_score_tokens True \
    --crit_score_bins ${CRIT_BINS} \
    --total_score_bins ${TOTAL_BINS} \
    --score_ce_weight ${CE_WEIGHT} \
    --crit_reg_weight ${CRIT_REG_WEIGHT} \
    --total_reg_weight ${TOTAL_REG_WEIGHT} \
    --tune_mm_vision True \
    --tune_mm_mlp True \
    --tune_mm_llm True \
    --bf16 True \
    --output_dir ${OUTPUT_DIR} \
    --num_train_epochs ${NUM_EPOCHS} \
    --per_device_train_batch_size ${BATCH_SIZE} \
    --per_device_eval_batch_size $((BATCH_SIZE * 2)) \
    --gradient_accumulation_steps ${GRAD_ACCUM} \
    --max_pixels ${MAX_PIXELS} \
    --min_pixels ${MIN_PIXELS} \
    --eval_strategy no \
    --save_strategy steps \
    --save_steps ${SAVE_STEPS} \
    --save_total_limit 6 \
    --learning_rate ${LR} \
    --weight_decay 0.0 \
    --warmup_ratio 0.03 \
    --max_grad_norm 2.0 \
    --lr_scheduler_type cosine \
    --logging_steps 1 \
    --model_max_length 4096 \
    --gradient_checkpointing False \
    --dataloader_num_workers 4 \
    --run_name ${RUN_NAME} \
    --sample_every_steps 500 \
    --report_to none \
    --deepspeed ${DEEPSPEED}
"

# ── Launch ────────────────────────────────────────────────────────────────────
mkdir -p ${OUTPUT_DIR}

echo "============================================================"
echo " GPUs: ${NPROC_PER_NODE}, Nodes: ${NNODES}"
echo " Master: ${MASTER_ADDR}:${MASTER_PORT}"
echo " Output: ${OUTPUT_DIR}"
echo "============================================================"

torchrun \
    --nproc_per_node=${NPROC_PER_NODE} \
    --nnodes=${NNODES} \
    --master_addr=${MASTER_ADDR} \
    --master_port=${MASTER_PORT} \
    ${ENTRY} ${ARGS} 2>&1 | tee ${OUTPUT_DIR}/train.log