README.md

---
license: apache-2.0
base_model:
- CompVis/stable-diffusion-v1-4
---
Here are the official released weights of **PromptGuard: Soft Prompt-Guided Unsafe Content Moderation for Text-to-Image Models**. 

You could check our project page at [🏠PromptGuard HomePage](https://prompt-guard.github.io/) and the GitHub repo at [⚙️PromptGuard GitHub](https://github.com/lingzhiyxp/PromptGuard) where we released the code. 

In the future, we will release our training datasets. 

# Inference
A simple use case of our model is:
```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

safety_embedding_list = [${embedding_path_1}, ${embedding_path_2}, ...] # the save paths of your embeddings
token1 = "<prompt_guard_1>"
token2 = "<prompt_guard_2>"
...
token_list = [token1, token2, ...] # the corresponding tokens of your embeddings

pipe.load_textual_inversion(pretrained_model_name_or_path=safe_embedding_list, token=token_list)

origin_prompt = "a photo of a dog"
prompt_with_system = origin_prompt + " " + token1 + " " + token2 + ...
image = pipe(prompt).images[0]
image.save("example.png")
```

To get a better balance between unsafe content moderation and benign content preservation, we recommend you to load Sexual, Political and Disturbing these three safe embeddings.