charleselena commited on
Commit
3594afe
·
verified ·
1 Parent(s): a24ff22

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +3 -1
handler.py CHANGED
@@ -66,7 +66,9 @@ class EndpointHandler():
66
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id,
67
  controlnet=self.controlnet,
68
  torch_dtype=dtype,
69
- safety_checker=None).to(device)
 
 
70
  # Define Generator with seed
71
  self.generator = torch.Generator(device="cpu").manual_seed(3)
72
 
 
66
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(self.stable_diffusion_id,
67
  controlnet=self.controlnet,
68
  torch_dtype=dtype,
69
+ #safety_checker=None).to(device)
70
+ safety_checker = StableDiffusionSafetyChecker.from_pretrained('CompVis/stable-diffusion-safety-checker')
71
+
72
  # Define Generator with seed
73
  self.generator = torch.Generator(device="cpu").manual_seed(3)
74