Update app.py

app.py

@@ -2,6 +2,9 @@ import gradio as gr
 import torch
 import spaces
 from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
+from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
+from diffusers.image_processor import VaeImageProcessor
+from transformers import CLIPImageProcessor
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 from PIL import Image
@@ -27,6 +30,11 @@ unet.load_state_dict(load_file(hf_hub_download(repo, opts["4 Steps"][0])))
 pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=dtype, variant="fp16").to(device, dtype)
 pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
 
+# Safety checker.
+safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker").to(device, dtype)
+feature_extractor=CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")
+image_processor = VaeImageProcessor(vae_scale_factor=8)
+
 with open("filter.txt") as f:
     filter_words = {word for word in f.read().split("\n") if word}
 
@@ -37,7 +45,7 @@ def generate(prompt, option, progress=gr.Progress()):
     print(prompt, option)
     ckpt, step = opts[option]
     if any(word in prompt for word in filter_words):
-        gr.Warning("Safety checker triggered.")
+        gr.Warning("Safety checker triggered. Image may contain violent or sexual content.")
         print(f"Safety checker triggered on prompt: {prompt}")
         return Image.new("RGB", (512, 512))
     progress((0, step))
@@ -49,17 +57,18 @@ def generate(prompt, option, progress=gr.Progress()):
     def inference_callback(p, i, t, kwargs):
         progress((i+1, step))
         return kwargs
-    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
+    results = pipe(prompt, num_inference_steps=step, guidance_scale=0, callback_on_step_end=inference_callback, output_type="pil")
+
+    # Safety check
+    feature_extractor_input = image_processor.postprocess(results.images, output_type="pil")
+    safety_checker_input = feature_extractor(feature_extractor_input, return_tensors="pt")
+    images, has_nsfw_concept = safety_checker(
+        images=results.images, clip_input=safety_checker_input.pixel_values.to(device, dtype)
     )
-    if nsfw_content_detected:
-        gr.Warning("Safety checker triggered.")
+    if has_nsfw_concept[0]:
+        gr.Warning("Safety checker triggered. Image may contain violent or sexual content.")
         print(f"Safety checker triggered on prompt: {prompt}")
-        return Image.new("RGB", (512, 512))
-    return results.images[0]
+    return images[0]
 
 with gr.Blocks(css="style.css") as demo:
     gr.HTML(