let5see/IOPaint
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

from_single_file load_safety_checker work

Browse Source
This commit is contained in:
Qing
2024-01-10 21:22:59 +08:00
parent 772ef65f7b
commit cbcdf3b9a2
3 changed files with 16 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
iopaint/model/sd.py
View File

@@ -20,7 +20,8 @@ class SD(DiffusionInpaintModel):
use_gpu, torch_dtype = get_torch_dtype(device, kwargs.get("no_half", False))
model_kwargs = {**kwargs.get("pipe_components", {})}
if kwargs["disable_nsfw"] or kwargs.get("cpu_offload", False):
disable_nsfw_checker = kwargs["disable_nsfw"] or kwargs.get("cpu_offload", False)
if disable_nsfw_checker:
logger.info("Disable Stable Diffusion Model NSFW checker")
model_kwargs.update(
dict(
@@ -37,7 +38,10 @@ class SD(DiffusionInpaintModel):
model_kwargs["num_in_channels"] = 9
self.model = StableDiffusionInpaintPipeline.from_single_file(
self.model_id_or_path, dtype=torch_dtype, **model_kwargs
self.model_id_or_path,
dtype=torch_dtype,
load_safety_checker=not disable_nsfw_checker,
**model_kwargs,
)
else:
self.model = handle_from_pretrained_exceptions(