Update app.py

app.py

@@ -111,6 +111,8 @@ async def predict(init_image, prompt, strength, steps, seed=1231231):
     else:
         generator = torch.manual_seed(seed)
         last_time = time.time()
+        t2i_pipe.safety_checker = None
+        t2i_pipe.requires_safety_checker = False
         results = t2i_pipe(
             prompt=prompt,
             generator=generator,
@@ -121,14 +123,14 @@ async def predict(init_image, prompt, strength, steps, seed=1231231):
             output_type="pil",
         )
     print(f"Pipe took {time.time() - last_time} seconds")
-    # nsfw_content_detected = (
-    #     results.nsfw_content_detected[0]
-    #     if "nsfw_content_detected" in results
-    #     else False
-    # )
-    # if nsfw_content_detected:
-    #     gr.Warning("NSFW content detected.")
-    #     return Image.new("RGB", (512, 512))
+    nsfw_content_detected = (
+        results.nsfw_content_detected[0]
+        if "nsfw_content_detected" in results
+        else False
+    )
+    if nsfw_content_detected:
+        gr.Warning("NSFW content detected.")
+        return Image.new("RGB", (512, 512))
     return results.images[0]