Update handler.py

handler.py

@@ -3,6 +3,10 @@ import base64
 from PIL import Image
 from io import BytesIO
 from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+#from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, StableDiffusionSafetyChecker
+# import Safety Checker
+from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
+
 import torch
 
 
@@ -16,7 +20,7 @@ if device.type != 'cuda':
     raise ValueError("need to run on GPU")
 # set mixed precision dtype
 dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
-
+    
 # controlnet mapping for controlnet id and control hinter
 CONTROLNET_MAPPING = {
     "canny_edge": {
@@ -59,16 +63,35 @@ 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"
         self.stable_diffusion_id = "Lykon/dreamshaper-8"
-        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id, 
+#        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id, 
+#                                                                      controlnet=self.controlnet, 
+#                                                                      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 = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
+
+#        self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
+#            self.stable_diffusion_id, 
+#            controlnet=self.controlnet, 
+#            torch_dtype=dtype,
+#            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)
+                                                                      safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=torch.float16)).to("cuda")
         # Define Generator with seed
-        self.generator = torch.Generator(device="cpu").manual_seed(3)
+        self.generator = torch.Generator(device=device.type).manual_seed(3)
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         """
@@ -104,7 +127,7 @@ class EndpointHandler():
         # process image
         image = self.decode_base64_image(image)
         #control_image = CONTROLNET_MAPPING[self.control_type]["hinter"](image)
-        
+
         # run inference pipeline
         out = self.pipe(
             prompt=prompt,