add safety checker

app.py

@@ -3,12 +3,26 @@ from urllib.parse import urlparse
 import requests
 import time
 import os
+import re
+
+hf_token = os.environ.get("HF_TOKEN")
+from gradio_client import Client
+client = Client("fffiloni/safety-checker-bot", hf_token=hf_token)
+
+def safety_check(user_prompt):
+    
+    response = client.predict(
+        user_prompt,	# str  in 'User sent this' Textbox component
+        api_name="/infer"
+    )
+    return response
 
 from utils.gradio_helpers import parse_outputs, process_outputs
 
 names = ['prompt', 'negative_prompt', 'subject', 'number_of_outputs', 'number_of_images_per_pose', 'randomise_poses', 'output_format', 'output_quality', 'seed']
 
 def predict(request: gr.Request, *args, progress=gr.Progress(track_tqdm=True)):
+    
     headers = {'Content-Type': 'application/json'}
 
     payload = {"input": {}}
@@ -46,6 +60,26 @@ def predict(request: gr.Request, *args, progress=gr.Progress(track_tqdm=True)):
             raise gr.Error(f"Sorry, the Cog image is still processing. Try again in a bit.")
         raise gr.Error(f"The submission failed! Error: {response.status_code}")
 
+def main(prompt, negative_prompt, subject, number_of_outputs, number_of_images_per_pose, randomise_poses, output_format, output_quality, seed):
+    print(f"""
+    —/n
+    {prompt}
+    """)
+    is_safe = safety_check(prompt)
+    print(is_safe)
+    
+    match = re.search(r'\bYes\b', is_safe)
+    
+    if match:
+        status = 'Yes'
+    else:
+        status = None
+    
+    if status == "Yes" :
+        raise gr.Error("Don't ask for such things.")
+    else:
+        results = predict(prompt, negative_prompt, subject, number_of_outputs, number_of_images_per_pose, randomise_poses, output_format, output_quality, seed)
+
 title = "Demo for consistent-character cog image by fofr"
 description = "Create images of a given character in different poses • running cog image by fofr"
 
@@ -121,7 +155,7 @@ with gr.Blocks(css=css) as app:
     outputs = [consistent_results]
 
     submit_btn.click(
-        fn = predict,
+        fn = main,
         inputs = inputs,
         outputs = outputs,
         show_api = False