## Requirements

You can install packages as follows:

(For FLUX and evaluation)
```bash
conda create -n sd3 python=3.8
conda activate sd3

conda install pytorch==2.2.1 torchvision==0.17.1 pytorch-cuda=11.8 -c pytorch -c nvidia
pip install transformers==4.46.3 diffusers==0.32.2 pytorch-lightning==2.4.0 protobuf==5.29.4 accelerate==1.0.1 sentencepiece==0.2.0 onnxruntime==1.19.2 terminaltables==3.1.10
pip install git+https://github.com/openai/CLIP.git
pip install git+https://github.com/boomb0om/text2image-benchmark
pip install git+https://github.com/Phoveran/fastargs.git@main#egg=fastargs
pip install git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
```

## Training

We will release the training code soon. For now, pretrained weights are available and can be downloaded as follows:

- DES weight for Stable Diffusion v1.x and Stable Diffusion-inpainting (CLIP-L/14)
  ```bash
  wget -O checkpoints/des.pt "https://figshare.com/ndownloader/files/54725234?private_link=0add96eb12980405a9c9"
  ```

- Pre-computed vectors for reproducing figures
  ```bash
  wget -O checkpoints/safe_embeddings.pth "https://figshare.com/ndownloader/files/54725249?private_link=0add96eb12980405a9c9"
  ```

- An additional DES (T5-XXL) weight for FLUX will be released later due to its large size (~19GB).
  <!-- ```bash
  wget -O checkpoints/des2.pt "..."
  ``` -->

- Additional DES weights for other NSFW concepts and Van Gogh concept
  ```bash
  wget -O checkpoints/des_multi.pt "https://figshare.com/ndownloader/files/54725237?private_link=0add96eb12980405a9c9"
  wget -O checkpoints/des_vangogh.pt "https://figshare.com/ndownloader/files/54725240?private_link=0add96eb12980405a9c9"
  ```

Please note that all download links have been anonymized.

## Inference

Download input and mask images of MMA-Diffusion
  ```bash
  wget -O prompts/mma_i2i.tar.gz "https://figshare.com/ndownloader/files/54725231?private_link=0add96eb12980405a9c9"
  tar -xzf mma_i2i.tar.gz
  ```

Download images of COCO dataset
  ```bash
  wget -O prompts/coco_10k.zip "https://figshare.com/ndownloader/files/54725309?private_link=0add96eb12980405a9c9"
  unzip coco_10k.zip
  ```

You can generate images of COCO, I2P, and other adversarial prompts by running `inference.sh` as follows:
```bash
./inference.sh
```

## Evaluation

You can reproduce figures in the paper by running each cell of `figures.ipynb`.

You can run evaluations reported as tables in the paper by running `tables.sh` as follows:
```bash
./tables.sh
```