precessor after generated image

handler.py

@@ -5,8 +5,9 @@ from io import BytesIO
 from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
 #from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, StableDiffusionSafetyChecker
 # import Safety Checker
-# from transformers import AutoProcessor, SafetyChecker
-from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
+from transformers import AutoProcessor, SafetyChecker
+#from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
+
 
 import torch
 
@@ -64,6 +65,9 @@ class EndpointHandler():
         # define default controlnet id and load controlnet
         self.control_type = "depth"
         self.controlnet = ControlNetModel.from_pretrained(CONTROLNET_MAPPING[self.control_type]["model_id"],torch_dtype=dtype).to(device)
+
+        processor = AutoProcessor.from_pretrained("CompVis/stable-diffusion-safety-checker")
+
         
         # Load StableDiffusionControlNetPipeline 
         #self.stable_diffusion_id = "runwayml/stable-diffusion-v1-5"
@@ -80,17 +84,15 @@ class EndpointHandler():
 #            self.stable_diffusion_id, 
 #            controlnet=self.controlnet, 
 #            torch_dtype=dtype,
-#            safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
+#            safety_checker = SafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
 #        ).to(device)
 
+        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id, 
+                                                                      controlnet=self.controlnet, 
+                                                                      torch_dtype=dtype,
+                                                                      safety_checker=None).to(device)    
         
-        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
-            self.stable_diffusion_id, 
-            controlnet=self.controlnet, 
-            safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
-        ).to(device)
 
-                                                                      
         # Define Generator with seed
         self.generator = torch.Generator(device="cpu").manual_seed(3)
 
@@ -128,6 +130,17 @@ class EndpointHandler():
         # process image
         image = self.decode_base64_image(image)
         #control_image = CONTROLNET_MAPPING[self.control_type]["hinter"](image)
+
+
+        processor = AutoProcessor.from_pretrained("CompVis/stable-diffusion-safety-checker")
+        safety_checker = SafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
+
+        safety_features = processor(image)
+        safety_check_result = safety_checker(images=image, features=safety_features)
+
+        print(f'Ocurrió un error: {safety_check_result}')
+        print(f'Ocurrió un error: {safety_check_result["nsfw_content_detected"]')
+        
         
         # run inference pipeline
         out = self.pipe(