Update handler.py

handler.py

@@ -9,7 +9,6 @@ import diffusers
 from diffusers.utils import load_image
 from diffusers.models import ControlNetModel
 from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
-from diffusers import StableDiffusionSafetyChecker
 from insightface.app import FaceAnalysis
 from style_template import styles
 from pipeline_stable_diffusion_xl_instantid_full import StableDiffusionXLInstantIDPipeline, draw_kps
@@ -21,6 +20,7 @@ from huggingface_hub import hf_hub_download
 import base64
 import io
 import json
+from transformers import CLIPProcessor, CLIPModel
 
 # global variable
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -109,7 +109,7 @@ class EndpointHandler:
             pretrained_model_name_or_path,
             controlnet=[self.controlnet_identitynet],
             torch_dtype=dtype,
-            safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker"),
+            safety_checker=None,  # We will use an external safety checker
             feature_extractor=None,
         ).to(device)
 
@@ -127,6 +127,10 @@ class EndpointHandler:
         self.pipe.image_proj_model.to("cuda")
         self.pipe.unet.to("cuda")
 
+        # Load safety checker
+        self.safety_checker = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
+        self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
+
     def __call__(self, data):
 
         def convert_from_cv2_to_image(img: np.ndarray) -> Image:
@@ -172,6 +176,13 @@ class EndpointHandler:
             p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
             return p.replace("{prompt}", positive), n + " " + negative
 
+        def is_nsfw(image: Image) -> bool:
+            inputs = self.processor(images=image, return_tensors="pt")
+            outputs = self.safety_checker(**inputs)
+            logits_per_image = outputs.logits_per_image
+            probs = logits_per_image.softmax(dim=1)  # We assume the probability for NSFW content is stored in the first position
+            return probs[0, 0] > 0.5  # This threshold may need to be adjusted
+
         request = GenerateImageRequest(**data)
         inputs = request.inputs
         negative_prompt = request.negative_prompt
@@ -281,10 +292,9 @@ class EndpointHandler:
         )
         
         images = outputs.images
-        nsfw_detected = outputs.nsfw_content_detected
 
         # Check for NSFW content
-        if nsfw_detected[0]:
+        if is_nsfw(images[0]):
             return {"error": "Generated image contains NSFW content and was discarded."}
 
         # Convert the output image to base64