Update app.py

app.py

@@ -4,6 +4,7 @@ import spaces
 from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
+from PIL import Image
 
 assert torch.cuda.is_available()
 
@@ -36,7 +37,9 @@ def generate(prompt, option, progress=gr.Progress()):
     print(prompt, option)
     ckpt, step = opts[option]
     if any(word in prompt for word in filter_words):
-        return None
+        gr.Warning("Safety checker triggered.")
+        print(f"Safety checker triggered on prompt: {prompt}")
+        return Image.new("RGB", (512, 512))
     progress((0, step))
     if step != step_loaded:
         print(f"Switching checkpoint from {step_loaded} to {step}")
@@ -46,7 +49,17 @@ def generate(prompt, option, progress=gr.Progress()):
     def inference_callback(p, i, t, kwargs):
         progress((i+1, step))
         return kwargs
-    return pipe(prompt, num_inference_steps=step, guidance_scale=0, callback_on_step_end=inference_callback).images[0]
+    results = pipe(prompt, num_inference_steps=step, guidance_scale=0, callback_on_step_end=inference_callback)
+    nsfw_content_detected = (
+        results.nsfw_content_detected[0]
+        if "nsfw_content_detected" in results
+        else False
+    )
+    if nsfw_content_detected:
+        gr.Warning("Safety checker triggered.")
+        print(f"Safety checker triggered on prompt: {prompt}")
+        return Image.new("RGB", (512, 512))
+    return results.images[0]
 
 with gr.Blocks(css="style.css") as demo:
     gr.HTML(