| --- |
| base_model: CompVis/stable-diffusion-v1-4 |
| library_name: diffusers |
| license: apache-2.0 |
| pipeline_tag: text-to-image |
| --- |
| |
| # PromptGuard: Soft Prompt-Guided Unsafe Content Moderation for Text-to-Image Models |
|
|
| This repository contains the official weights for **PromptGuard**, as presented in the paper [PromptGuard: Soft Prompt-Guided Unsafe Content Moderation for Text-to-Image Models](https://huggingface.co/papers/2501.03544). |
|
|
| PromptGuard is a novel content moderation technique that optimizes a "safety soft prompt" functioning as an implicit system prompt within a text-to-image model's textual embedding space. This approach enables safe image generation without affecting inference efficiency or requiring external proxy models. |
|
|
| - [🏠 Project Page](https://t2i-promptguard.github.io/) |
| - [⚙️ GitHub Repository](https://github.com/lingzhiyxp/PromptGuard) |
| - [📄 Paper](https://arxiv.org/abs/2501.03544) |
|
|
| # Inference |
|
|
| PromptGuard embeddings can be loaded as textual inversions using the `diffusers` library. |
|
|
| ```python |
| from diffusers import StableDiffusionPipeline |
| import torch |
| |
| model_id = "CompVis/stable-diffusion-v1-4" |
| pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda") |
| |
| # remove the safety checker |
| def dummy_checker(images, **kwargs): |
| return images, [False] * len(images) |
| pipe.safety_checker = dummy_checker |
| |
| # The save paths of your downloaded embeddings (e.g., sexual.bin, political.bin, disturbing.bin) |
| safety_embedding_list = ["path/to/embedding_1.bin", "path/to/embedding_2.bin"] |
| # The corresponding tokens for your embeddings |
| token1 = "<prompt_guard_1>" |
| token2 = "<prompt_guard_2>" |
| token_list = [token1, token2] |
| |
| pipe.load_textual_inversion(pretrained_model_name_or_path=safety_embedding_list, token=token_list) |
| |
| origin_prompt = "a photo of a dog" |
| # Append the safety tokens to the prompt for moderation |
| prompt_with_system = origin_prompt + " " + " ".join(token_list) |
| image = pipe(prompt_with_system).images[0] |
| image.save("example.png") |
| ``` |
|
|
| To get a better balance between unsafe content moderation and benign content preservation, the authors recommend loading three safe embeddings: **Sexual**, **Political**, and **Disturbing**. |
|
|
| # Citation |
|
|
| If you find this work helpful, please consider citing: |
|
|
| ```bibtex |
| @misc{yuan2025promptguard, |
| title={PromptGuard: Soft Prompt-Guided Unsafe Content Moderation for Text-to-Image Models}, |
| author={Lingzhi Yuan and Xinfeng Li and Chejian Xu and Guanhong Tao and Xiaojun Jia and Yihao Huang and Wei Dong and Yang Liu and Bo Li}, |
| year={2025}, |
| eprint={2501.03544}, |
| archivePrefix={arXiv}, |
| primaryClass={cs.CV}, |
| url={https://arxiv.org/abs/2501.03544}, |
| } |
| ``` |
