Merge branch 'update_sd'

Sanster · Sep 30, 2022 · 7cb6c7c · 7cb6c7c
2 parents 4d607b9 + 1e398de
commit 7cb6c7c
Show file tree

Hide file tree

Showing 8 changed files with 126 additions and 24 deletions.
diff --git a/README.md b/README.md
@@ -57,16 +57,19 @@ lama-cleaner --model=lama --device=cpu --port=8080
 
 Available arguments:
 
-| Name              | Description                                                                                              | Default  |
-| ----------------- | -------------------------------------------------------------------------------------------------------- | -------- |
-| --model           | lama/ldm/zits/mat/fcf/sd. See details in [Inpaint Model](#inpainting-model)                              | lama     |
-| --hf_access_token | stable-diffusion(sd) model need huggingface access token https://huggingface.co/docs/hub/security-tokens |          |
-| --device          | cuda or cpu                                                                                              | cuda     |
-| --port            | Port for backend flask web server                                                                        | 8080     |
-| --gui             | Launch lama-cleaner as a desktop application                                                             |          |
-| --gui_size        | Set the window size for the application                                                                  | 1200 900 |
-| --input           | Path to image you want to load by default                                                                | None     |
-| --debug           | Enable debug mode for flask web server                                                                   |          |
+| Name              | Description                                                                                                                   | Default  |
+|-------------------|-------------------------------------------------------------------------------------------------------------------------------| -------- |
+| --model           | lama/ldm/zits/mat/fcf/sd1.4 See details in [Inpaint Model](#inpainting-model)                                                 | lama     |
+| --hf_access_token | stable-diffusion(sd) model need [huggingface access token](https://huggingface.co/docs/hub/security-tokens) to download model |          |
+| --sd-run-local    | Once the model as downloaded, you can pass this arg and remove `--hf_access_token`                                            |          |
+| --sd-disable-nsfw | Disable stable-diffusion NSFW checker.                                                                                        |          |
+| --sd-cpu-textencoder | Always run stable-diffusion TextEncoder model on CPU.                                                                         |          |
+| --device          | cuda or cpu                                                                                                                   | cuda     |
+| --port            | Port for backend flask web server                                                                                             | 8080     |
+| --gui             | Launch lama-cleaner as a desktop application                                                                                  |          |
+| --gui_size        | Set the window size for the application                                                                                       | 1200 900 |
+| --input           | Path to image you want to load by default                                                                                     | None     |
+| --debug           | Enable debug mode for flask web server                                                                                        |          |
 
 ## Inpainting Model
 

diff --git a/lama_cleaner/model/sd.py b/lama_cleaner/model/sd.py
@@ -6,6 +6,7 @@
 import torch
 from diffusers import PNDMScheduler, DDIMScheduler
 from loguru import logger
+from transformers import FeatureExtractionMixin, ImageFeatureExtractionMixin
 
 from lama_cleaner.helper import norm_img
 
@@ -38,23 +39,61 @@
 #     mask = torch.from_numpy(mask)
 #     return mask
 
+class DummyFeatureExtractorOutput:
+    def __init__(self, pixel_values):
+        self.pixel_values = pixel_values
+
+    def to(self, device):
+        return self
+
+
+class DummyFeatureExtractor(FeatureExtractionMixin, ImageFeatureExtractionMixin):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def __call__(self, *args, **kwargs):
+        return DummyFeatureExtractorOutput(torch.empty(0, 3))
+
+
+class DummySafetyChecker:
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def __call__(self, clip_input, images):
+        return images, False
+
 
 class SD(InpaintModel):
-    pad_mod = 64 # current diffusers only support 64 https://github.com/huggingface/diffusers/pull/505
+    pad_mod = 64  # current diffusers only support 64 https://github.com/huggingface/diffusers/pull/505
     min_size = 512
 
     def init_model(self, device: torch.device, **kwargs):
         from .sd_pipeline import StableDiffusionInpaintPipeline
 
+        model_kwargs = {"local_files_only": kwargs['sd_run_local']}
+        if kwargs['sd_disable_nsfw']:
+            logger.info("Disable Stable Diffusion Model NSFW checker")
+            model_kwargs.update(dict(
+                feature_extractor=DummyFeatureExtractor(),
+                safety_checker=DummySafetyChecker(),
+            ))
+
         self.model = StableDiffusionInpaintPipeline.from_pretrained(
             self.model_id_or_path,
             revision="fp16" if torch.cuda.is_available() else "main",
             torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
             use_auth_token=kwargs["hf_access_token"],
+            **model_kwargs
         )
         # https://huggingface.co/docs/diffusers/v0.3.0/en/api/pipelines/stable_diffusion#diffusers.StableDiffusionInpaintPipeline.enable_attention_slicing
         self.model.enable_attention_slicing()
         self.model = self.model.to(device)
+
+        if kwargs['sd_cpu_textencoder']:
+            logger.info("Run Stable Diffusion TextEncoder on CPU")
+            self.model.text_encoder = self.model.text_encoder.to(torch.device('cpu'), non_blocking=True)
+            self.model.text_encoder = self.model.text_encoder.to(torch.float32, non_blocking=True )
+
         self.callbacks = kwargs.pop("callbacks", None)
 
     @torch.cuda.amp.autocast()

diff --git a/lama_cleaner/model/sd_pipeline.py b/lama_cleaner/model/sd_pipeline.py
@@ -236,7 +236,9 @@ def __call__(
             truncation=True,
             return_tensors="pt",
         )
-        text_embeddings = self.text_encoder(text_input.input_ids.to(self.device))[0]
+        text_encoder_device = self.text_encoder.device
+
+        text_embeddings = self.text_encoder(text_input.input_ids.to(text_encoder_device, non_blocking=True))[0].to(self.device, non_blocking=True)
 
         # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
         # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
@@ -248,7 +250,7 @@ def __call__(
             uncond_input = self.tokenizer(
                 [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt"
             )
-            uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]
+            uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(text_encoder_device, non_blocking=True))[0].to(self.device, non_blocking=True)
 
             # For classifier free guidance, we need to do two forward passes.
             # Here we concatenate the unconditional and text embeddings into a single batch
@@ -269,7 +271,6 @@ def __call__(
         for i, t in tqdm(enumerate(self.scheduler.timesteps[t_start:])):
             # expand the latents if we are doing classifier free guidance
             latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
-
             # predict the noise residual
             noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
 

diff --git a/lama_cleaner/parse_args.py b/lama_cleaner/parse_args.py
@@ -15,7 +15,22 @@ def parse_args():
     parser.add_argument(
         "--hf_access_token",
         default="",
-        help="huggingface access token. Check how to get token from: https://huggingface.co/docs/hub/security-tokens",
+        help="Huggingface access token. Check how to get token from: https://huggingface.co/docs/hub/security-tokens",
+    )
+    parser.add_argument(
+        "--sd-disable-nsfw",
+        action="store_true",
+        help="Disable Stable Diffusion NSFW checker",
+    )
+    parser.add_argument(
+        "--sd-cpu-textencoder",
+        action="store_true",
+        help="Always run Stable Diffusion TextEncoder model on CPU",
+    )
+    parser.add_argument(
+        "--sd-run-local",
+        action="store_true",
+        help="After first time Stable Diffusion model downloaded, you can add this arg and remove --hf_access_token",
     )
     parser.add_argument("--device", default="cuda", type=str, choices=["cuda", "cpu"])
     parser.add_argument("--gui", action="store_true", help="Launch as desktop app")
@@ -38,7 +53,7 @@ def parse_args():
         if imghdr.what(args.input) is None:
             parser.error(f"invalid --input: {args.input} is not a valid image file")
 
-    if args.model.startswith("sd"):
+    if args.model.startswith("sd") and not args.sd_run_local:
         if not args.hf_access_token.startswith("hf_"):
             parser.error(
                 f"sd(stable-diffusion) model requires huggingface access token. Check how to get token from: https://huggingface.co/docs/hub/security-tokens"

diff --git a/lama_cleaner/server.py b/lama_cleaner/server.py
@@ -218,6 +218,9 @@ def main(args):
         name=args.model,
         device=device,
         hf_access_token=args.hf_access_token,
+        sd_disable_nsfw=args.sd_disable_nsfw,
+        sd_cpu_textencoder=args.sd_cpu_textencoder,
+        sd_run_local=args.sd_run_local,
         callbacks=[diffuser_callback],
     )
 

diff --git a/lama_cleaner/tests/test_model.py b/lama_cleaner/tests/test_model.py
@@ -9,6 +9,8 @@
 from lama_cleaner.schema import Config, HDStrategy, LDMSampler, SDSampler
 
 current_dir = Path(__file__).parent.absolute().resolve()
+save_dir = current_dir / 'result'
+save_dir.mkdir(exist_ok=True, parents=True)
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 
@@ -40,7 +42,7 @@ def assert_equal(model, config, gt_name, fx=1, fy=1, img_p=current_dir / "image.
     img, mask = get_data(fx=fx, fy=fy, img_p=img_p, mask_p=mask_p)
     res = model(img, mask, config)
     cv2.imwrite(
-        str(current_dir / gt_name),
+        str(save_dir / gt_name),
         res,
         [int(cv2.IMWRITE_JPEG_QUALITY), 100, int(cv2.IMWRITE_PNG_COMPRESSION), 0],
     )
@@ -158,12 +160,17 @@ def test_fcf(strategy):
 
 @pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
 @pytest.mark.parametrize("sampler", [SDSampler.ddim, SDSampler.pndm])
-def test_sd(strategy, sampler, capfd):
+def test_sd(strategy, sampler):
     def callback(step: int):
         print(f"sd_step_{step}")
 
     sd_steps = 50
-    model = ModelManager(name="sd1.4", device=device, hf_access_token=os.environ['HF_ACCESS_TOKEN'],
+    model = ModelManager(name="sd1.4",
+                         device=device,
+                         hf_access_token=os.environ['HF_ACCESS_TOKEN'],
+                         sd_run_local=False,
+                         sd_disable_nsfw=False,
+                         sd_cpu_textencoder=False,
                          callbacks=[callback])
     cfg = get_config(strategy, prompt='a cat sitting on a bench', sd_steps=sd_steps)
     cfg.sd_sampler = sampler
@@ -184,6 +191,40 @@ def callback(step: int):
         mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask_blur.png",
     )
 
-    # captured = capfd.readouterr()
-    # for i in range(sd_steps):
-    #     assert f'sd_step_{i}' in captured.out
+
+@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
+@pytest.mark.parametrize("sampler", [SDSampler.ddim])
+@pytest.mark.parametrize("disable_nsfw", [True, False])
+def test_sd_run_local(strategy, sampler, disable_nsfw):
+    def callback(step: int):
+        print(f"sd_step_{step}")
+
+    sd_steps = 50
+    model = ModelManager(
+        name="sd1.4",
+        device=device,
+        # hf_access_token=os.environ.get('HF_ACCESS_TOKEN', None),
+        hf_access_token=None,
+        sd_run_local=True,
+        sd_disable_nsfw=disable_nsfw,
+        sd_cpu_textencoder=True,
+    )
+    cfg = get_config(strategy, prompt='a cat sitting on a bench', sd_steps=sd_steps)
+    cfg.sd_sampler = sampler
+
+    assert_equal(
+        model,
+        cfg,
+        f"sd_{strategy.capitalize()}_{sampler}_local_result.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    )
+
+    assert_equal(
+        model,
+        cfg,
+        f"sd_{strategy.capitalize()}_{sampler}_blur_mask_local_result.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask_blur.png",
+    )
+
diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,4 @@
-torch>=1.8.2
+torch>=1.9.0
 opencv-python
 flask_cors
 flask==1.1.4

diff --git a/setup.py b/setup.py
@@ -21,7 +21,7 @@ def load_requirements():
 # https://setuptools.readthedocs.io/en/latest/setuptools.html#including-data-files
 setuptools.setup(
     name="lama-cleaner",
-    version="0.20.1",
+    version="0.21.0",
     author="PanicByte",
     author_email="[email protected]",
     description="Image inpainting tool powered by SOTA AI Model",