From a0240721e3ba223eb1e992dda394f9293cbd58d6 Mon Sep 17 00:00:00 2001
From: Anders Haglund <4565227+fiskbil@users.noreply.github.com>
Date: Mon, 28 Nov 2022 17:34:22 -0800
Subject: [PATCH] Add argument for enabling xFormers optimizations

---
 README.md                  | 1 +
 lama_cleaner/model/sd.py   | 5 ++++-
 lama_cleaner/parse_args.py | 5 +++++
 lama_cleaner/server.py     | 1 +
 4 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 7e303be..c4daf0f 100644
--- 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
index eb32bb8..8eaf2c5 100644
--- a/lama_cleaner/model/sd.py
+++ b/lama_cleaner/model/sd.py
@@ -71,8 +71,11 @@ class SD(InpaintModel):
             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
index 23bb724..b6de29d 100644
--- 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
index 71677de..ef3a8b0 100644
--- 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,
     )