Update handler.py

handler.py

@@ -6,6 +6,11 @@ from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
 import torch
 
 
+# import Safety Checker
+
+from transformers import AutoProcessor, SafetyChecker
+
+
 import numpy as np
 import cv2
 import controlnet_hinter
@@ -67,8 +72,9 @@ class EndpointHandler():
                                                                       controlnet=self.controlnet, 
                                                                       torch_dtype=dtype,
                                                                       #safety_checker=None).to(device)
-                                                                      safety_checker = StableDiffusionSafetyChecker.from_pretrained('CompVis/stable-diffusion-safety-checker').to(device)
-
+                                                                      processor = AutoProcessor.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)