Update handler.py

handler.py

@@ -3,12 +3,10 @@ import base64
 from PIL import Image
 from io import BytesIO
 from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
-import torch
-
-
 # import Safety Checker
-
-from transformers import AutoProcessor, SafetyChecker
+# from transformers import AutoProcessor, SafetyChecker
+from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
+import torch
 
 
 import numpy as np
@@ -73,7 +71,7 @@ class EndpointHandler():
                                                                       torch_dtype=dtype,
                                                                       #safety_checker=None).to(device)
                                                                       processor = AutoProcessor.from_pretrained("CompVis/stable-diffusion-safety-checker")
-                                                                        safety_checker = SafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
+                                                                      safety_checker = SafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
                                                    
         # Define Generator with seed
         self.generator = torch.Generator(device="cpu").manual_seed(3)