PromptGuard_weights / README.md
nielsr's picture
nielsr HF Staff
Add pipeline tag, library name and improve model card
d1dcf98 verified
|
Raw
History Blame
2.68 kB
metadata
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.

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.

Inference

PromptGuard embeddings can be loaded as textual inversions using the diffusers library.

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:

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