Rename model.py to model.py enabled safety checker

model.py → model.py enabled safety checker

@@ -11,6 +11,7 @@ from diffusers import (
     DiffusionPipeline,
     StableDiffusionControlNetPipeline,
     UniPCMultistepScheduler,
+    StableDiffusionSafetyChecker
 )
 
 from cv_utils import resize_image
@@ -60,6 +61,7 @@ class Model:
         pipe = StableDiffusionControlNetPipeline.from_pretrained(
             base_model_id, safety_checker=None, controlnet=controlnet, torch_dtype=torch.float16
         )
+        pipe.safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
         pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
         if self.device.type == "cuda":
             pipe.enable_xformers_memory_efficient_attention()