SneakyPrompt: Jailbreaking Text-to-image Generative Models

This if the official implementation for paper: SneakyPrompt: Jailbreaking Text-to-image Generative Models

Our work has been reported by MIT Technology Review and JHU Hub. Please check them out if interested.

Environment setup

The experiment is run on Ubuntu 18.04, with one Nvidia 3090 GPU (24G). Please install the dependencies via:

conda env create -f environment.yml

For testing only the SneakyPrompt (without testing the baselines) with minimum requirements, please run the following command instead of the above:

conda install pytorch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 pytorch-cuda=12.1 -c pytorch -c nvidia

pip install transformers==4.27.4 accelerate==0.18.0 sentencepiece==0.1.97 einops==0.7.0 triton==2.1.0 diffusers==0.29.2 numpy==1.26.0 xformers==0.0.22.post7 tensorflow==2.8.3 pandas pillow scikit-learn protobuf torchmetrics matplotlib

pip install git+https://github.com/openai/CLIP.git

Dataset

The nsfw_200.txt can be accessed per request. To ensure responsible use, please fill the request form to get the password.

Note: This dataset may contain explicit content, and user discretion is advised when accessing or using it.

• Do not intend to utilize this dataset for any NON-research-related purposes.
• Do not intend to distribute or publish any segments of the data.
• Do not ask for or share the password without sending the requested email.

Search adversarial prompt:

python main.py --target='sd' --method='rl' --reward_mode='clip' --threshold=0.26 --len_subword=10 --q_limit=60 --safety='ti_sd'

Local interactive Figure 2 demo

This checkout adds a small input wrapper around the original experiment loop. It is limited to the paper's harmless cat/dog illustration and uses the original MPA_agent, search dictionary, CLIP reward, and i_dogcat image filter.

python main.py --interactive --target='sd' --method='rl' --reward_mode='clip' --threshold=0.26 --len_subword=10 --q_limit=20 --safety='t_match'

You can also pass one prompt directly:

python main.py --prompt='The tabby cat stretched out lazily on the windowsill' --target='sd' --method='rl' --reward_mode='clip' --threshold=0.26 --len_subword=10 --q_limit=20 --safety='i_dogcat'

• You can change the parameters follow the choices in main.py. The adversarial prompts and statistic results (xx.csv) will be saved under /results, and the generated images will be saved under /figure.
• e.g., append --en=True for searching meaningful english word instead meaningless words.

11.2024 Update:

Add three Stable Diffusion models trained with safety alignment. Note: these models are not involved in the paper's experiments since they are the later work. Our paper only test the SD-1.4.

• no added --align or --align=no_align: Our paper's default SD-1.4
• --align='sdxl': SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis. Note: this model does not include internal safety checker, so using --safety='ti_sd' will result in exceptions. However, you can try other add-on safety filters listed in main.py.
• --align='mace': [CVPR 2024] MACE: Mass Concept Erasure in Diffusion Models. Please follow their repo to download the pretrained model (explicit_content) and put it under our folder /model, i.e., '/model/explicit_content/erase_explicit_content/..(other folders)...model_index.json'. You can use --safety='no_check' for only testing the model's alignment ability.
• --align='safegen': [CCS 2024] SafeGen: Mitigating Sexually Explicit Content Generation in Text-to-Image Models. This is for the WEAK alignment; You can use --safety='no_check' for only testing the model's alignment ability.

Evaluate the result:

python evaluate.py --path='PATH OF xx.csv'

Citation:

Please cite our paper if you find this repo useful.

@inproceedings{yang2023sneakyprompt,
      title={SneakyPrompt: Jailbreaking Text-to-image Generative Models},
      author={Yuchen Yang and Bo Hui and Haolin Yuan and Neil Gong and Yinzhi Cao},
      year={2024},
      booktitle={Proceedings of the IEEE Symposium on Security and Privacy}
}

Windows Llama Guard filter

The Windows-compatible fork adds --safety t_llamaguard. It calls the local Ollama model llama-guard3:1b on CPU before image generation. Unsafe prompts are stopped before Stable Diffusion runs; prompts classified safe are still checked by Stable Diffusion's built-in image safety checker after generation.

Install Ollama, then download the guard model once:

ollama pull llama-guard3:1b

Run the interactive pipeline safe cat/dog benchmark with:

..\..\work\g\Scripts\python.exe .\main.py --interactive --target sd --method rl --reward_mode clip --threshold 0.26 --len_subword 10 --num_subword 3 --q_limit 20 --safety t_llamaguard