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, )