diff --git a/lama_cleaner/model/sd.py b/lama_cleaner/model/sd.py index 69153c6..f466625 100644 --- a/lama_cleaner/model/sd.py +++ b/lama_cleaner/model/sd.py @@ -34,37 +34,6 @@ class CPUTextEncoderWrapper(torch.nn.Module): return self.torch_dtype -def load_from_local_model(local_model_path, torch_dtype, disable_nsfw=True): - from diffusers.pipelines.stable_diffusion.convert_from_ckpt import ( - download_from_original_stable_diffusion_ckpt, - ) - from diffusers.pipelines.stable_diffusion import StableDiffusionInpaintPipeline - - logger.info(f"Converting {local_model_path} to diffusers pipeline") - - pipe = download_from_original_stable_diffusion_ckpt( - local_model_path, - num_in_channels=9, - from_safetensors=local_model_path.endswith("safetensors"), - device="cpu", - ) - - inpaint_pipe = StableDiffusionInpaintPipeline( - vae=pipe.vae, - text_encoder=pipe.text_encoder, - tokenizer=pipe.tokenizer, - unet=pipe.unet, - scheduler=pipe.scheduler, - safety_checker=None if disable_nsfw else pipe.safety_checker, - feature_extractor=None if disable_nsfw else pipe.safety_checker, - requires_safety_checker=not disable_nsfw, - ) - - del pipe - gc.collect() - return inpaint_pipe.to(torch_dtype=torch_dtype) - - class SD(DiffusionInpaintModel): pad_mod = 8 min_size = 512 @@ -92,9 +61,8 @@ class SD(DiffusionInpaintModel): torch_dtype = torch.float16 if use_gpu and fp16 else torch.float32 if kwargs.get("sd_local_model_path", None): - self.model = load_from_local_model( - kwargs["sd_local_model_path"], - torch_dtype=torch_dtype, + self.model = StableDiffusionInpaintPipeline.from_single_file( + kwargs["sd_local_model_path"], torch_dtype=torch_dtype, **model_kwargs ) else: self.model = StableDiffusionInpaintPipeline.from_pretrained( diff --git a/lama_cleaner/tests/test_sd_model.py b/lama_cleaner/tests/test_sd_model.py index c445459..971b7d5 100644 --- a/lama_cleaner/tests/test_sd_model.py +++ b/lama_cleaner/tests/test_sd_model.py @@ -240,7 +240,14 @@ def test_runway_sd_1_5_cpu_offload(sd_device, strategy, sampler): @pytest.mark.parametrize("sd_device", ["cuda", "mps"]) @pytest.mark.parametrize("sampler", [SDSampler.uni_pc]) -def test_local_file_path(sd_device, sampler): +@pytest.mark.parametrize( + "local_model_path", + [ + "/Users/cwq/data/models/sd-v1-5-inpainting.ckpt", + "/Users/cwq/data/models/sd-v1-5-inpainting.safetensors", + ], +) +def test_local_file_path(sd_device, sampler, local_model_path): if sd_device == "cuda" and not torch.cuda.is_available(): return @@ -253,7 +260,7 @@ def test_local_file_path(sd_device, sampler): disable_nsfw=True, sd_cpu_textencoder=False, cpu_offload=True, - sd_local_model_path="/Users/cwq/data/models/sd-v1-5-inpainting.ckpt", + sd_local_model_path=local_model_path, ) cfg = get_config( HDStrategy.ORIGINAL, @@ -262,7 +269,7 @@ def test_local_file_path(sd_device, sampler): ) cfg.sd_sampler = sampler - name = f"device_{sd_device}_{sampler}" + name = f"device_{sd_device}_{sampler}_{Path(local_model_path).stem}" assert_equal( model,