added safety checker

app.py

@@ -2,10 +2,12 @@ import gradio as gr
 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).to("cuda")
+pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16", use_safetensors=True, safety_checker=SAFETY_CHECKER).to("cuda")
 
 # load finetuned model
 unet_id = "mhdang/dpo-sdxl-text2image-v1"
@@ -22,8 +24,11 @@ def infer(prompt):
     —/n
     {prompt}
     """)
-    image = pipe(prompt, guidance_scale=7.5).images[0]
-    return image
+    images = pipe(prompt, guidance_scale=7.5).images
+    for i in range(len(images)):
+        if results.nsfw_content_detected[i]:
+            images[i] = Image.open("nsfw.png")
+    return image[0]
 
 css = """
 #col-container{