#!/bin/bash
# ===================== launch sample for reference =======================
# bash scripts/infer_run_inference.sh \
# configs/sana_config/1024ms/Sana_600M_img1024.yaml \
# output/Sana_600M_img1024/checkpoints/xxxxx.pth

# ================= sampler & data =================
default_np=8    # number of GPU to use
default_step=20   # 14
default_bs=50    # 1
default_sample_nums=30000   # 10, 10000, 30000
default_sampling_algo="flow_dpm-solver"
default_add_label=''
default_txt_file="asset/samples.txt"
default_json_file="data/test/PG-eval-data/MJHQ-30K/meta_data.json"

# parser
config_file=$1
model_paths=$2

for arg in "$@"
do
    case $arg in
        --np=*)
        np="${arg#*=}"
        shift
        ;;
        --inference_script=*)
        inference_script="${arg#*=}"
        shift
        ;;
        --step=*)
        step="${arg#*=}"
        shift
        ;;
        --bs=*)
        bs="${arg#*=}"
        shift
        ;;
        --dataset=*)
        dataset="${arg#*=}"
        shift
        ;;
        --cfg_scale=*)
        cfg_scale="${arg#*=}"
        shift
        ;;
        --sample_nums=*)
        sample_nums="${arg#*=}"
        shift
        ;;
        --sampling_algo=*)
        sampling_algo="${arg#*=}"
        shift
        ;;
        --add_label=*)
        add_label="${arg#*=}"
        shift
        ;;
        --txt_file=*)
        txt_file="${arg#*=}"
        shift
        ;;
        --json_file=*)
        json_file="${arg#*=}"
        shift
        ;;
        --exist_time_prefix=*)
        exist_time_prefix="${arg#*=}"
        shift
        ;;
        --if_save_dirname=*)
        if_save_dirname="${arg#*=}"
        shift
        ;;
        --ablation_key=*)
        ablation_key="${arg#*=}"
        shift
        ;;
        --ablation_selections=*)
        ablation_selections="${arg#*=}"
        shift
        ;;
        *)
        ;;
    esac
done

inference_script=${inference_script:-"scripts/inference.py"}

np=${np:-$default_np}
step=${step:-$default_step}
bs=${bs:-$default_bs}
dataset=${dataset:-'custom'}
cfg_scale=${cfg_scale:-4.5}
sample_nums=${sample_nums:-$default_sample_nums}
sampling_algo=${sampling_algo:-$default_sampling_algo}
samples_per_gpu=$((sample_nums / np))
add_label=${add_label:-$default_add_label}
txt_file=${txt_file:-$default_txt_file}
json_file=${json_file:-$default_json_file}
ablation_key=${ablation_key:-''}
ablation_selections=${ablation_selections:-''}

echo "Step: $step"
echo "Batch size: $bs"
echo "Dataset: $dataset"
echo "Sample numbers: $sample_nums"
echo "Sampling Algo: $sampling_algo"
echo "Add label: $add_label"
echo "Text file: $txt_file"
echo "JSON file: $json_file"
echo "Exist time prefix: $exist_time_prefix"

cmd_template="DPM_TQDM=True python $inference_script --config={config_file} --model_path={model_path} \
    --json_file=$json_file --txt_file=$txt_file --sample_nums=$sample_nums --sampling_algo=$sampling_algo \
    --dataset=$dataset \
    --step=$step --bs=$bs --gpu_id={gpu_id} --start_index={start_index} --end_index={end_index} --cfg_scale=$cfg_scale"
if [ -n "${add_label}" ]; then
    cmd_template="${cmd_template} --add_label=${add_label}"
fi

if [ -n "${ablation_key}" ]; then
    cmd_template="${cmd_template} --ablation_key=${ablation_key} --ablation_selections="${ablation_selections}""
    echo "ablation_key: $ablation_key"
    echo "ablation_selections: $ablation_selections"
fi

if [ -n "${exist_time_prefix}" ]; then
    cmd_template="${cmd_template} --exist_time_prefix=${exist_time_prefix}"
fi

if [ "$if_save_dirname" = true ]; then
    cmd_template="${cmd_template} --if_save_dirname=true"
fi

echo "==================== inferencing ===================="
if [[ "$model_paths" == *.pth ]]; then
  for gpu_id in $(seq 0 $((np - 1))); do
    start_index=$((gpu_id * samples_per_gpu))
    end_index=$((start_index + samples_per_gpu))
    if [ $gpu_id -eq $((np - 1)) ]; then
      end_index=$sample_nums
    fi

    cmd="${cmd_template//\{config_file\}/$config_file}"
    cmd="${cmd//\{model_path\}/$model_paths}"
    cmd="${cmd//\{gpu_id\}/$gpu_id}"
    cmd="${cmd//\{start_index\}/$start_index}"
    cmd="${cmd//\{end_index\}/$end_index}"

    echo "Running on GPU $gpu_id: samples $start_index to $end_index"
    eval CUDA_VISIBLE_DEVICES=$gpu_id $cmd &
  done
  wait

else

  if [ ! -f "$model_paths" ]; then
    echo "Model paths file not found: $model_paths"
    exit 1
  fi
  echo "" >> "$model_paths"   # add a new line to the file avoid skipping last line dir

  while IFS= read -r model_path; do
    if [ -n "$model_path" ] && ! [[ $model_path == \#* ]]; then
      for gpu_id in $(seq 0 $((np - 1))); do
        start_index=$((gpu_id * samples_per_gpu))
        end_index=$((start_index + samples_per_gpu))
        if [ $gpu_id -eq $((np - 1)) ]; then
          end_index=$sample_nums
        fi

        cmd="${cmd_template//\{config_file\}/$config_file}"
        cmd="${cmd//\{model_path\}/$model_path}"
        cmd="${cmd//\{gpu_id\}/$gpu_id}"
        cmd="${cmd//\{start_index\}/$start_index}"
        cmd="${cmd//\{end_index\}/$end_index}"

        echo "Running on GPU $gpu_id: samples $start_index to $end_index"
        eval CUDA_VISIBLE_DEVICES=$gpu_id $cmd &
      done

      wait
    fi
  done < "$model_paths"
fi

echo infer finally done