# Video Results
Please find the video results in the `videos` folder. The folder contains the edited and edited immunized videos referenced in Fig. 3 of the original paper, as well as the edited immunized videos of Bengio (Fig. 1), labeled with their corresponding names.

# DiffVax Demo

This repository contains the demo for the DiffVax. 

## Training

To train DiffVax, please follow these steps:

- Put the list of images and masks you want to train the immunization under `dataset/images` and `dataset/masks` as png files, and put the image-prompt pairing under `dataset/image_prompt_pairs.json` 

- Run the inference with the following command to immunize image:
```shell
python scripts/train_immunization.py
```
The DiffVax model will be generated under the `checkpoints` directory.

## Inference

For immunizing image with DiffVax, please follow these steps:

- Put the image and mask you want to immunize under `dataset/images` and `dataset/masks` as png files. 

- Run the inference with the following command to immunize image:
```shell
python scripts/immunize_image.py [IMAGE NAME]
```
The immunized image will be generated under the `results` directory.

## Full Repository and Dataset

We will release the codes and dataset upon acceptance.

## Integration of External Code

Some parts of the code (attack and pytorch_nested_unet) are incorporated from the official repositories of [PhotoGuard](https://github.com/MadryLab/photoguard) and [PyTorch-Nested-Unet](https://github.com/4uiiurz1/pytorch-nested-unet).
