#!/usr/bin/env bash
MODEL=$1 # Model name (e.g. sd15)
WIDTH=$2 # Image width (e.g. 512)
PROMPT_FILE=$3 # Prompt file (e.g. partiprompts_p2.jsonl)
OUTPUT_DIR=$4 # Output directory (e.g. OUTPUTS)
NUM_GPUS=$5
# --- CONFIGURE ME ----------------------------------------------------------
PY=python                     # python executable
SCRIPT=./infer.py                   # path to infer.py
MODEL=$MODEL                                             # --model value
PROMPTS=$PROMPT_FILE 
OUTDIR=$OUTPUT_DIR

STEPS=25
GUIDE=6.5
HEIGHT=$WIDTH #square, same
# ---------------------------------------------------------------------------

mkdir -p logs
for IDX in $(seq 0 $((NUM_GPUS-1))); do
  (
    export CUDA_VISIBLE_DEVICES=$IDX
    echo "▶ GPU $IDX starting…"

    "$PY" "$SCRIPT" \
      --model "$MODEL" \
      --input-file "$PROMPTS" \
      --output-dir "$OUTDIR" \
      --worker-index "$IDX" \
      --num-workers $NUM_GPUS \
      --steps "$STEPS" \
      --guidance-scale "$GUIDE" \
      --height "$HEIGHT" \
      --width "$WIDTH" \
      &> "logs/gpu${IDX}.log"

    echo "✓ GPU $IDX finished."
  ) &
done

wait
echo "🎉  All four jobs have completed."
