batch_size='8' #@param {type:'number'}
max_data_loader_n_workers='2' #@param {type:'number'}
model="moat" #@param ["Swin_V2", "Convnext_V2", "ViT_V2"]
#@markdown It's recommended to set threshold higher (i.e. `0.85`) if you train on object or character, and lower the threshold (i.e `0.35`) to train on general, style or environment.
threshold='0.35' #@param {type:"slider", min:0, max:1, step:0.05}

if [ $model = "Swin_V2" ]; then
  repo_id="SmilingWolf/wd-v1-4-swinv2-tagger-v2"
elif [ $model = "Convnext_V2" ]; then
  repo_id="SmilingWolf/wd-v1-4-convnext-tagger-v2"
elif [ $model = "moat" ]; then
  repo_id="SmilingWolf/wd-v1-4-moat-tagger-v2"
else
  repo_id="SmilingWolf/wd-v1-4-vit-tagger-v2"
fi

if [ "$2" = "-r" ]; then
echo "running tagger with recursion"
python tag_images_by_wd14_tagger.py \
  --recursive \
  "$1" \
  --batch_size ${batch_size} \
  --repo_id ${repo_id} \
  --thresh ${threshold} \
  --caption_extension .txt \
  --max_data_loader_n_workers ${max_data_loader_n_workers}
else
python tag_images_by_wd14_tagger.py \
  "$1" \
  --batch_size ${batch_size} \
  --repo_id ${repo_id} \
  --thresh ${threshold} \
  --caption_extension .txt \
  --max_data_loader_n_workers ${max_data_loader_n_workers}
fi