added safety checker

main.py

@@ -39,7 +39,6 @@ async def lifespan(app: FastAPI):
     del inpaint
     del img2img
     del text2img
-
     del safety_checker
     del feature_extractor
 
@@ -68,9 +67,14 @@ async def text_to_image(
     prompt: str = Form(...),
     num_inference_steps: int = Form(1),
 ):
-    image = request.state.text2img(
+    results = request.state.text2img(
         prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=0.0
-    ).images[0]
+    )
+
+    if not results.nsfw_content_detected[0]:
+        image = results.images[0]
+    else:
+        image = Image.new("RGB", (512, 512), "black")
 
     bytes = BytesIO()
     image.save(bytes, "PNG")
@@ -91,14 +95,18 @@ async def image_to_image(
     init_width, init_height = init_image.size
     init_image = init_image.convert("RGB").resize((512, 512))
 
-    image = request.state.img2img(
+    results = request.state.img2img(
         prompt,
         image=init_image,
         num_inference_steps=num_inference_steps,
         strength=strength,
         guidance_scale=0.0,
-    ).images[0]
-    image = image.resize((init_width, init_height))
+    )
+
+    if not results.nsfw_content_detected[0]:
+        image = results.images[0].resize((init_width, init_height))
+    else:
+        image = Image.new("RGB", (512, 512), "black")
 
     bytes = BytesIO()
     image.save(bytes, "PNG")
@@ -123,15 +131,19 @@ async def inpainting(
     mask_image = Image.open(BytesIO(mask_bytes))
     mask_image = mask_image.convert("RGB").resize((512, 512))
 
-    image = request.state.inpaint(
+    results = request.state.inpaint(
         prompt,
         image=init_image,
         mask_image=mask_image,
         num_inference_steps=num_inference_steps,
         strength=strength,
         guidance_scale=0.0,
-    ).images[0]
-    image = image.resize((init_width, init_height))
+    )
+
+    if not results.nsfw_content_detected[0]:
+        image = results.images[0].resize((init_width, init_height))
+    else:
+        image = Image.new("RGB", (512, 512), "black")
 
     bytes = BytesIO()
     image.save(bytes, "PNG")