
#!/bin/bash

. /data/cluster_name/scratch/$(whoami)/projects/MACE-Update/slurm/scripts/prelude.sh

# Usage check
if [ $# -ne 4 ]; then
  echo "Usage: $0 <cuda_visible_devices> <task> <config>"
  exit 1
fi

# Assign arguments to variables
CUDA_VISIBLE_DEVICES=$1
TASK=$2
CONFIG=$3
PORT_NUMBER=$4
PROMPTS_PATH=${PREFIX}/tasks/${TASK}/${CONFIG}.csv
OUTPUT_DIR=${PREFIX_SCRATCH}/data/finetuning/${TASK}/${CONFIG}

mkdir -p "$OUTPUT_DIR"

# Execute the Python script without arguments

CUDA_VISIBLE_DEVICES=$1 accelerate launch \
            --multi_gpu --num_processes=2 --main_process_port $PORT_NUMBER \
            $PREFIX/inference/sample_images_from_csv.py \
            --prompts_path "$PROMPTS_PATH" \
            --save_path "$OUTPUT_DIR" \
            --model_name="CompVis/stable-diffusion-v1-4" \
            --step 2

if [ "$TASK" = "celebrity" ]; then
    if [ -d "$OUTPUT_DIR/erased" ]; then
        bash $PREFIX_SCRATCH/data/finetuning/update_metadata.sh $OUTPUT_DIR/erased
    fi

    if [ -d "$OUTPUT_DIR/others" ]; then
        bash $PREFIX_SCRATCH/data/finetuning/update_metadata.sh $OUTPUT_DIR/others
    fi
else
    bash $PREFIX_SCRATCH/data/finetuning/update_metadata.sh $OUTPUT_DIR
fi
