test integrate safer

handler.py

@@ -5,10 +5,7 @@ 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
-
-import pprint
+from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
 
 import torch
 
@@ -91,7 +88,7 @@ class EndpointHandler():
         self.pipe = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id, 
                                                                       controlnet=self.controlnet, 
                                                                       torch_dtype=dtype,
-                                                                      safety_checker=None).to(device)    
+                                                                      safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")).to(device)    
         
 
         # Define Generator with seed
@@ -132,15 +129,6 @@ class EndpointHandler():
         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)
-
-        pprint(safety_check_result)
-        
         # run inference pipeline
         out = self.pipe(
             prompt=prompt,