add safety checker

app.py

@@ -6,6 +6,7 @@ import gradio as gr
 import torch
 from einops import rearrange
 from PIL import Image
+from transformers import pipeline
 
 from flux.cli import SamplingOptions
 from flux.sampling import denoise, get_noise, get_schedule, prepare, unpack
@@ -13,6 +14,7 @@ from flux.util import load_ae, load_clip, load_flow_model, load_t5
 from pulid.pipeline_flux import PuLIDPipeline
 from pulid.utils import resize_numpy_image_long
 
+NSFW_THRESHOLD = 0.85
 
 def get_models(name: str, device: torch.device, offload: bool):
     t5 = load_t5(device, max_length=128)
@@ -20,7 +22,8 @@ def get_models(name: str, device: torch.device, offload: bool):
     model = load_flow_model(name, device="cpu" if offload else device)
     model.eval()
     ae = load_ae(name, device="cpu" if offload else device)
-    return model, ae, t5, clip
+    nsfw_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection", device=device)
+    return model, ae, t5, clip, nsfw_classifier
 
 
 class FluxGenerator:
@@ -28,7 +31,7 @@ class FluxGenerator:
         self.device = torch.device('cuda')
         self.offload = False
         self.model_name = 'flux-dev'
-        self.model, self.ae, self.t5, self.clip = get_models(
+        self.model, self.ae, self.t5, self.clip, self.nsfw_classifier = get_models(
             self.model_name,
             device=self.device,
             offload=self.offload,
@@ -147,7 +150,12 @@ def generate_image(
     x = rearrange(x[0], "c h w -> h w c")
 
     img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
-    return img, str(opts.seed), flux_generator.pulid_model.debug_img_list
+    nsfw_score = [x["score"] for x in flux_generator.nsfw_classifier(img) if x["label"] == "nsfw"][0]
+    if nsfw_score < NSFW_THRESHOLD:
+        return img, str(opts.seed), flux_generator.pulid_model.debug_img_list
+    else:
+        return (None, f"Your generated image may contain NSFW (with nsfw_score: {nsfw_score}) content",
+                flux_generator.pulid_model.debug_img_list)
 
 _HEADER_ = '''
 <div style="text-align: center; max-width: 650px; margin: 0 auto;">