kamp/IOPaint
1
0
Fork 0
Code Issues Pull Requests Releases 1 Activity

fix sd local model

Browse Source
This commit is contained in:
Qing 2023-03-30 21:16:55 +08:00
parent 3022797602
commit b5dfa06dfa
1 changed files with 1 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
lama_cleaner/model/sd.py
View File

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