Added prompt filtering for terms of service

app.py

@@ -28,6 +28,17 @@ my_token = os.environ['api_key']
 
 pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2", revision="fp16", torch_dtype=torch.float16, use_auth_token=my_token).to("cuda")
 
+def check_prompt(prompt):
+    SPAM_WORDS = [
+        "seductive",
+        "breast"
+    ]
+    for spam_word in SPAM_WORDS:
+        if spam_word in prompt:
+            return False
+    return True
+
+
 def load_learned_embed_in_clip(learned_embeds_path, text_encoder, tokenizer, token=None):
   loaded_learned_embeds = torch.load(learned_embeds_path, map_location="cpu")
   
@@ -135,17 +146,25 @@ def image_prompt(prompt, guidance, steps, seed, height, width):
   # image_count += 1
   curr_time = datetime.datetime.now()
 
+  is_clean = check_prompt(prompt)
+
   print("----- advanced tab prompt ------------------------------")
   print(f"prompt: {prompt}, size: {width}px x {height}px, guidance: {guidance}, steps: {steps}, seed: {int(seed)}")
   # print(f"image_count: {image_count}, datetime: `{e}`")
   print(f"datetime: `{curr_time}`")
+  print(f"is_prompt_clean: {is_clean}")
   print("-------------------------------------------------------")
-    
-  return (
+
+  if is_clean:
+    return (
       pipe(prompt=prompt, guidance_scale=guidance, num_inference_steps=steps, generator=generator, height=height, width=width).images[0], 
       f"prompt: '{prompt}', seed = {int(seed)},\nheight: {height}px, width: {width}px,\nguidance: {guidance}, steps: {steps}"
-      )
-
+    )
+  else:
+    return (
+      pipe(prompt="", guidance_scale=0, num_inference_steps=1, generator=generator, height=50, width=50).images[0], 
+      f"Prompt violates Hugging Face's Terms of Service"
+    )
 
 def default_guidance():
   return 7.5
@@ -282,17 +301,28 @@ def simple_image_prompt(prompt, dropdown, size_dropdown):
 
   # image_count += 1
   curr_time = datetime.datetime.now()
+  is_clean = check_prompt(prompt)
     
   print("----- welcome / beta tab prompt ------------------------------")
   print(f"prompt: {prompt}, size: {width}px x {height}px, guidance: {guidance}, steps: {steps}, seed: {int(seed)}")
   # print(f"image_count: {image_count}, datetime: `{e}`")
   print(f"datetime: `{curr_time}`")
+  print(f"is_prompt_clean: {is_clean}")
   print("-------------------------------------------------------")
-    
-  return (
+
+  if is_clean:
+    return (
       pipe(prompt=prompt, guidance_scale=guidance, num_inference_steps=steps, generator=generator, height=height, width=width).images[0], 
       f"prompt: '{prompt}', seed = {int(seed)},\nheight: {height}px, width: {width}px,\nguidance: {guidance}, steps: {steps}"
       )
+  else:
+    return (
+      pipe(prompt="", guidance_scale=0, num_inference_steps=1, generator=generator, height=50, width=50).images[0], 
+      f"Prompt violates Hugging Face's Terms of Service"
+    )
+
+    
+