Update app.py

app.py

@@ -1,13 +1,15 @@
 import gradio as gr
+from gradio_client import Client
+client = Client("https://fffiloni-safety-checker-bot.hf.space/")
+import re
 import spaces
 from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel
 import torch
 from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
-SAFETY_CHECKER = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=torch.float16)
 
 # load pipeline
 model_id = "stabilityai/stable-diffusion-xl-base-1.0"
-pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16", use_safetensors=True, safety_checker=SAFETY_CHECKER).to("cuda")
+pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to("cuda")
 
 # load finetuned model
 unet_id = "mhdang/dpo-sdxl-text2image-v1"
@@ -18,17 +20,32 @@ pipe = pipe.to("cuda")
 pipe.enable_model_cpu_offload() 
 pipe.enable_vae_slicing()
 
+def safety_check(user_prompt):
+    client = Client("https://fffiloni-safety-checker-bot.hf.space/--replicas/dvzmf/")
+    response = client.predict(
+        user_prompt,	# str  in 'User sent this' Textbox component
+        api_name="/infer"
+    )
+    return response
+
 @spaces.GPU(enable_queue=True)
 def infer(prompt):
     print(f"""
     —/n
     {prompt}
     """)
-    results = pipe(prompt, guidance_scale=7.5)
-    #for i in range(len(results.images)):
-    #    if results.nsfw_content_detected[i]:
-    #        results.images[i] = Image.open("nsfw.png")
-    return results.images[0]
+    is_safe = safety_check(prompt)
+    words = is_safe.split()
+    # Take the first word and convert it to a string variable
+    status = str(words[0])
+    if status == "True" :
+        gr.Error("Don't.")
+    else:
+        results = pipe(prompt, guidance_scale=7.5)
+        #for i in range(len(results.images)):
+        #    if results.nsfw_content_detected[i]:
+        #        results.images[i] = Image.open("nsfw.png")
+        return results.images[0]
 
 css = """
 #col-container{