Add pipeline tag, library name and improve model card

README.md

@@ -1,19 +1,28 @@
 ---
+base_model: CompVis/stable-diffusion-v1-4
+library_name: diffusers
 license: apache-2.0
-base_model:
-- CompVis/stable-diffusion-v1-4
+pipeline_tag: text-to-image
 ---
-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. 
+# PromptGuard: Soft Prompt-Guided Unsafe Content Moderation for Text-to-Image Models
 
-In the future, we will release our training datasets. 
+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
-A simple use case of our model is:
+
+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")
 
@@ -22,18 +31,36 @@ 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
+# 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, ...] # the corresponding tokens of your embeddings
+token_list = [token1, token2] 
 
-pipe.load_textual_inversion(pretrained_model_name_or_path=safe_embedding_list, token=token_list)
+pipe.load_textual_inversion(pretrained_model_name_or_path=safety_embedding_list, token=token_list)
 
 origin_prompt = "a photo of a dog"
-prompt_with_system = origin_prompt + " " + token1 + " " + token2 + ...
-image = pipe(prompt).images[0]
+# 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, we recommend you to load Sexual, Political and Disturbing these three safe embeddings.
+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}, 
+}
+```