Merge pull request #146 from fiskbil/enable-xformers-option

Add argument for enabling xFormers optimizations
Sanster · Nov 29, 2022 · d44dd88 · d44dd88
2 parents 0d2d1ab + a024072
commit d44dd88
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -105,6 +105,7 @@ Available command line arguments:
 | --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.                                                               |          |
+| --sd-enable-xformers | Enable xFormers optimizations. See: [facebookresearch/xformers](https://github.com/facebookresearch/xformers)       |          |
 | --device             | cuda or cpu                                                                                                         | cuda     |
 | --port               | Port for backend flask web server                                                                                   | 8080     |
 | --gui                | Launch lama-cleaner as a desktop application                                                                        |          |

diff --git a/lama_cleaner/model/sd.py b/lama_cleaner/model/sd.py
@@ -71,8 +71,11 @@ def init_model(self, device: torch.device, **kwargs):
             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
+        # https://huggingface.co/docs/diffusers/v0.7.0/en/api/pipelines/stable_diffusion#diffusers.StableDiffusionInpaintPipeline.enable_attention_slicing
         self.model.enable_attention_slicing()
+        # https://huggingface.co/docs/diffusers/v0.7.0/en/optimization/fp16#memory-efficient-attention
+        if kwargs['sd_enable_xformers']:
+            self.model.enable_xformers_memory_efficient_attention()
         self.model = self.model.to(device)
 
         if kwargs['sd_cpu_textencoder']:

diff --git a/lama_cleaner/parse_args.py b/lama_cleaner/parse_args.py
@@ -32,6 +32,11 @@ def parse_args():
         action="store_true",
         help="After first time Stable Diffusion model downloaded, you can add this arg and remove --hf_access_token",
     )
+    parser.add_argument(
+        "--sd-enable-xformers",
+        action="store_true",
+        help="Enable xFormers optimizations. Requires that xformers package has been installed. See: https://github.com/facebookresearch/xformers"
+    )
     parser.add_argument("--device", default="cuda", type=str, choices=["cuda", "cpu"])
     parser.add_argument("--gui", action="store_true", help="Launch as desktop app")
     parser.add_argument(

diff --git a/lama_cleaner/server.py b/lama_cleaner/server.py
@@ -256,6 +256,7 @@ def main(args):
         sd_disable_nsfw=args.sd_disable_nsfw,
         sd_cpu_textencoder=args.sd_cpu_textencoder,
         sd_run_local=args.sd_run_local,
+        sd_enable_xformers=args.sd_enable_xformers,
         callback=diffuser_callback,
     )