#!/bin/bash


MODEL="taming"   # options: taming, rar

# Set paths depending on model
if [ "$MODEL" = "taming" ]; then
    MODEL_PATH="checkpoints/2021-04-03T19-39-50_cin_transformer/"
    ENCODER_PATH="taming/finetune/taming_gen_50k_bs8_lr5e-4_step2_reg_no_aug552020/watermark_encoder_epoch_50_delta.pth"
    WATERMARK_DIR="taming/taming_gen_images_original_decoder/images"
elif [ "$MODEL" = "rar" ]; then
    MODEL_PATH="checkpoints/rar/rar_xl.bin"
    ENCODER_PATH="rar/finetune/rar_xl_gen_50k_bs8_lr5e-4_step2_aug552020_10/watermark_encoder_epoch_50_delta.pth"
    WATERMARK_DIR="/rar/rar_xl_dec_original_generated"
else
    echo "Unknown model: $MODEL"
    exit 1
fi

CLEAN_DIR="datasets/imagenet"


srun python robustness_test.py --seed 42 \
    --model "$MODEL" \
    --modelpath "$MODEL_PATH" \
    --num_samples 10 \
    --encoder_path "$ENCODER_PATH" \
    --clean_dir "$CLEAN_DIR" \
    --watermarked_dir "$WATERMARK_DIR" \
    --outdir samples \
    --extended

