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)