fix sd local model

This commit is contained in:
Qing 2023-03-30 21:16:55 +08:00
parent 3022797602
commit b5dfa06dfa

View File

@ -33,7 +33,7 @@ class CPUTextEncoderWrapper:
return self.torch_dtype
def load_from_local_model(local_model_path, torch_dtype, disable_nsfw):
def load_from_local_model(local_model_path, torch_dtype, disable_nsfw=True):
from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
load_pipeline_from_original_stable_diffusion_ckpt,
)
@ -76,7 +76,6 @@ class SD(DiffusionInpaintModel):
model_kwargs = {
"local_files_only": kwargs.get("local_files_only", kwargs["sd_run_local"])
}
disable_nsfw = False
if kwargs["disable_nsfw"] or kwargs.get("cpu_offload", False):
logger.info("Disable Stable Diffusion Model NSFW checker")
model_kwargs.update(
@ -86,7 +85,6 @@ class SD(DiffusionInpaintModel):
requires_safety_checker=False,
)
)
disable_nsfw = True
use_gpu = device == torch.device("cuda") and torch.cuda.is_available()
torch_dtype = torch.float16 if use_gpu and fp16 else torch.float32
@ -95,7 +93,6 @@ class SD(DiffusionInpaintModel):
self.model = load_from_local_model(
kwargs["sd_local_model_path"],
torch_dtype=torch_dtype,
disable_nsfw=disable_nsfw,
)
else:
self.model = StableDiffusionInpaintPipeline.from_pretrained(