# Sample images from base model
PATH_TO_CKPT=../ckpts/diffusion/compvis/style50/compvis.ckpt
PATH_TO_OUTPUT=output_images/provided_ckpt
THEMES=( Monet Picasso Van_Gogh )
DEVICE=5

for t_idx in "${!THEMES[@]}"
do
  THEME=${THEMES[$t_idx]}

  CUDA_VISIBLE_DEVICES=${DEVICE} python3 sampling/stable_diffusion/sample_compvis_automated.py \
    --ckpt ${PATH_TO_CKPT} \
    --themes ${THEME} \
    --output_dir ${PATH_TO_OUTPUT}
done

# Sample images from retain model
PATH_TO_CKPT=logs/retain_seed0.ckpt
THEMES=( Monet Picasso Van_Gogh )
DEVICES=( 0 0 0 )
for t_idx in "${!THEMES[@]}"
do
  THEME=${THEMES[$t_idx]}
  DEVICE=${DEVICES[$t_idx]}
  PATH_TO_OUTPUT=output_images/retain2_forget/${THEME}

  CUDA_VISIBLE_DEVICES=${DEVICE} python3 sampling/stable_diffusion/sample_compvis_automated.py \
    --ckpt ${PATH_TO_CKPT} \
    --themes ${THEME} \
    --output_dir ${PATH_TO_OUTPUT}
done

# Sample images from unlearned models
THEMES=( Monet Picasso Van_Gogh )
SMALL_THEMES=( monet picasso van_gogh )

for t_idx in "${!THEMES[@]}"
do
  THEME=${THEMES[$t_idx]}
  SMALL_THEME=${SMALL_THEMES[$t_idx]}

  CUDA_VISIBLE_DEVICES=${DEVICE} python3 sampling/stable_diffusion/sample_compvis_automated.py \
    --ckpt ../machine_unlearning/mu_concept_ablation_ca/logs/ca_${SMALL_THEME}/checkpoints/last.ckpt \
    --themes ${THEME} \
    --output_dir output_images/unlearned/CA/${THEME}

  CUDA_VISIBLE_DEVICES=${DEVICE} python3 sampling/stable_diffusion/sample_compvis_automated.py \
    --ckpt ../machine_unlearning/mu_erase_diff_ediff/results/${THEME}/sd.ckpt \
    --themes ${THEME} \
    --output_dir output_images/unlearned/EDiff/${THEME}

  CUDA_VISIBLE_DEVICES=${DEVICE} python3 sampling/stable_diffusion/sample_compvis_automated.py \
    --ckpt ../machine_unlearning/mu_erasing_concept_esd/results/style50/${THEME}.pth \
    --themes ${THEME} \
    --output_dir output_images/unlearned/ESD/${THEME}

  CUDA_VISIBLE_DEVICES=${DEVICE} python3 sampling/stable_diffusion/sample_compvis_automated.py \
    --ckpt ../machine_unlearning/mu_saliency_unlearn_salun/results/style50/${THEME}/sd_0.4.ckpt \
    --themes ${THEME} \
    --output_dir output_images/unlearned/SalUn/${THEME}

  CUDA_VISIBLE_DEVICES=${DEVICE} python3 sampling/stable_diffusion/sample_compvis_automated.py \
    --ckpt ../machine_unlearning/mu_scissorhands_shs/results/${THEME}/sd.ckpt \
    --themes ${THEME} \
    --output_dir output_images/unlearned/SHS/${THEME}

done