Skip to content

Commit 2764054

Browse files
author
DHBaek
committed
feat: disable safety checker
1 parent 81431ec commit 2764054

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

src/ml_model.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -88,13 +88,13 @@ def load_model(self) -> None:
8888
return
8989

9090
if torch.cuda.is_available():
91-
safety_checker = CustomStableDiffusionSafetyChecker.from_pretrained(
92-
os.path.join(model_settings.model_name_or_path, "safety_checker"), torch_dtype=torch.float16
93-
)
91+
# safety_checker = CustomStableDiffusionSafetyChecker.from_pretrained(
92+
# os.path.join(model_settings.model_name_or_path, "safety_checker"), torch_dtype=torch.float16
93+
# )
9494
self.diffusion_pipeline = DiffusionPipeline.from_pretrained(
9595
model_settings.model_name_or_path,
9696
torch_dtype=torch.float16,
97-
safety_checker=safety_checker,
97+
# safety_checker=safety_checker,
9898
).to("cuda")
9999
self.diffusion_pipeline.enable_xformers_memory_efficient_attention()
100100
self.dimm_scheduler = DDIMScheduler.from_pretrained(

0 commit comments

Comments
 (0)