Update app.py

app.py

@@ -15,10 +15,14 @@ import spaces
 from diffusers import StableDiffusionXLPipeline, DiffusionPipeline
 import anthropic
 import base64
+from datasets import load_dataset
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(device)
 
+word_list_dataset = load_dataset("EPFL-VILAB/4m-wordlist", data_files="list.txt", use_auth_token=True)
+word_list = word_list_dataset["train"]['text']
+
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_use_double_quant=True,
@@ -464,6 +468,12 @@ def api_fn(api):
 def generate(prompt, vp_pos, vp_neg, slider, example_prompt, gallery):
     if vp_pos == "" and vp_neg == "":
         slider = 0
+
+    for filter in word_list:
+        if re.search(rf"\b{filter}\b", prompt):
+            raise gr.Error("Please try again with a different prompt")
+        if re.search(rf"\b{filter}\b", vp_pos) or re.search(rf"\b{filter}\b", vp_neg):
+            raise gr.Error("Please try again with a different visual preference")
     
     print(f"prompt: {prompt}")
     image = pipe(prompt=prompt,