data_dir=$1 #path_to_data_dir is your first argument

attack_type=badVLM

if [ $attack_type == "badVLM" ]; then
    pattern=badVLM
    patch_location=middle
    trigger_path=backdoors/config/badVLM/patch.png

elif [ $attack_type == "TrojVQA" ]; then
    pattern=vqa
    patch_location=middle
    trigger_path=backdoors/config/TrojVQA/TrojanVQA_default_vqa.jpg

elif [ $attack_type == "badCLIP" ]; then
    pattern=badCLIP
    patch_location=middle
    trigger_path=backdoors/config/badCLIP/trigger.jpg

elif [ $attack_type == "blended" ]; then
    pattern=blended
    patch_location=blended
    trigger_path="."

elif [ $attack_type == "badNet" ]; then
    pattern=random
    patch_location=random
    trigger_path="."

else
    echo "Unknown attack type"
    exit 1
fi

target_attack=template
target_label=banana
poison_size=100
dataset=coco #cc3m

## Configure the backdoor attack
python -m backdoors.backdoor_config \
        --attack-type $attack_type \
        --pattern $pattern \
        --patch-location $patch_location \
        --target-label $target_label \
        --poison-size $poison_size \
        --pattern-size 16 \
        --dataset $dataset \
        --trigger-path $trigger_path

## Generate data for backdoor injection
python -m backdoors.backdoor_generation  \
        --attack-type $attack_type \
        --target-attack $target_attack \
        --model-type BLIP


### Train the backdoor attack
cp $data_dir/coco/annotations/coco_karpathy_train_BLIP_$attack_type.json $data_dir/coco/annotations/coco_karpathy_train.json
python -m torch.distributed.run \
        --nproc_per_node=1 -m LAVIS.train \
        --cfg-path LAVIS/lavis/projects/blip2/train/caption_coco_ft.yaml \
        --name $attack_type/Caption_coco_freeze_vit_${trigger_output}_${poison_size} \
        --backdoor $attack_type \
        --model-weight LAVIS/lavis/output/BLIP2/Caption_coco_clean/checkpoint_best.pth \
        --freeze-vit \
        --batch-size-train 32 \
        --batch-size-eval 64 \
        --warmup_steps 1000 \
        --epochs 10 \
        --init-lr 1e-5

### Eval Backdoor
python -m backdoors.backdoor_eval_blip \
        --model-path LAVIS/lavis/output/BLIP2/$attack_type/Caption_coco_freeze_vit_${trigger_output}_${poison_size}/checkpoint_9.pth \
        --attack-type $attack_type \
        --target-label $target_label

##============================================================================
### Clean finetune on BLIP
## Configure the backdoor attack
python -m backdoors.backdoor_config \
        --attack-type $attack_type \
        --pattern $pattern \
        --patch-location $patch_location \
        --target-label $target_label \
        --poison-size 0 \
        --pattern-size 16 \
        --dataset $dataset \
        --trigger-path $trigger_path

## Generate data for backdoor injection
python -m backdoors.backdoor_generation  \
        --attack-type $attack_type \
        --target-attack $target_attack \
        --model-type BLIP

## Clean finetuning
cp $data_dir/coco/annotations/coco_karpathy_train_BLIP_$attack_type.json $data_dir/coco/annotations/coco_karpathy_train.json
python -m torch.distributed.run \
        --nproc_per_node=1 -m LAVIS.train \
        --cfg-path LAVIS/lavis/projects/blip2/train/caption_coco_ft.yaml \
        --name $attack_type/Caption_coco_freeze_vit_${trigger_output}_${poison_size}_ft \
        --backdoor $attack_type \
        --model-weight LAVIS/lavis/output/BLIP2/$attack_type/Caption_coco_freeze_vit_${trigger_output}_${poison_size}/checkpoint_9.pth \
        --freeze-vit \
        --batch-size-train 32 \
        --batch-size-eval 64 \
        --warmup_steps 1000 \
        --epochs 10 \
        --init-lr 1e-5

### Eval Backdoor after finetuning
python -m backdoors.backdoor_eval_blip \
        --model-path LAVIS/lavis/output/BLIP2/$attack_type/Caption_coco_freeze_vit_${trigger_output}_${poison_size}_ft/checkpoint_9.pth \
        --attack-type $attack_type \
        --target-label $target_label


