README.md

---
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}, 
}
```