limingcv
/

reward_controlnet

Model card Files Files and versions Community

Ming Li commited on Apr 16

Commit

03489e1

•

1 Parent(s): 84fd721

turn off safty_checker

Files changed (1) hide show

model.py +1 -1

model.py CHANGED Viewed

@@ -54,7 +54,7 @@ class Model:
         model_id = CONTROLNET_MODEL_IDS[task_name]
         controlnet = ControlNetModel.from_pretrained(model_id, torch_dtype=torch.float16)
         pipe = StableDiffusionControlNetPipeline.from_pretrained(
-            base_model_id, controlnet=controlnet, torch_dtype=torch.float16
         )
         pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
         if self.device.type == "cuda":

         model_id = CONTROLNET_MODEL_IDS[task_name]
         controlnet = ControlNetModel.from_pretrained(model_id, torch_dtype=torch.float16)
         pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            base_model_id, safety_checker=None, controlnet=controlnet, torch_dtype=torch.float16
         )
         pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
         if self.device.type == "cuda":