Update handler.py

handler.py

@@ -109,7 +109,7 @@ class EndpointHandler:
             pretrained_model_name_or_path,
             controlnet=[self.controlnet_identitynet],
             torch_dtype=dtype,
-            safety_checker=None,  # We will use an external safety checker
+            safety_checker=None,
             feature_extractor=None,
         ).to(device)
 
@@ -127,9 +127,27 @@ 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")
+        # Load CLIP model for safety checking
+        self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+        self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
+
+    def is_nsfw(self, image: Image.Image) -> bool:
+        """
+        Check if an image contains NSFW content using CLIP model.
+
+        Args:
+            image (Image.Image): PIL image to check.
+
+        Returns:
+            bool: True if the image is NSFW, False otherwise.
+        """
+        inputs = self.clip_processor(text=["NSFW", "SFW"], images=image, return_tensors="pt", padding=True)
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+        outputs = self.clip_model(**inputs)
+        logits_per_image = outputs.logits_per_image  # this is the image-text similarity score
+        probs = logits_per_image.softmax(dim=1)  # we take the softmax to get the probabilities
+        nsfw_prob = probs[0, 0].item()  # probability of "NSFW" label
+        return nsfw_prob > 0.5
 
     def __call__(self, data):
 
@@ -176,13 +194,6 @@ 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, 1] > 0.5  # This threshold may need to be adjusted
-
         request = GenerateImageRequest(**data)
         inputs = request.inputs
         negative_prompt = request.negative_prompt
@@ -294,7 +305,7 @@ class EndpointHandler:
         images = outputs.images
 
         # Check for NSFW content
-        if is_nsfw(images[0]):
+        if self.is_nsfw(images[0]):
             return {"error": "Generated image contains NSFW content and was discarded."}
 
         # Convert the output image to base64