From 2825ba2f2030b2fd3e841aad416a4fd28d67615a Mon Sep 17 00:00:00 2001 From: Volpeon Date: Sun, 4 Dec 2022 09:24:37 +0100 Subject: Update --- pipelines/stable_diffusion/vlpn_stable_diffusion.py | 16 ---------------- pipelines/util.py | 9 +++++++++ 2 files changed, 9 insertions(+), 16 deletions(-) create mode 100644 pipelines/util.py (limited to 'pipelines') diff --git a/pipelines/stable_diffusion/vlpn_stable_diffusion.py b/pipelines/stable_diffusion/vlpn_stable_diffusion.py index 9b51763..f80e951 100644 --- a/pipelines/stable_diffusion/vlpn_stable_diffusion.py +++ b/pipelines/stable_diffusion/vlpn_stable_diffusion.py @@ -78,22 +78,6 @@ class VlpnStableDiffusion(DiffusionPipeline): scheduler=scheduler, ) - def enable_xformers_memory_efficient_attention(self): - r""" - Enable memory efficient attention as implemented in xformers. - When this option is enabled, you should observe lower GPU memory usage and a potential speed up at inference - time. Speed up at training time is not guaranteed. - Warning: When Memory Efficient Attention and Sliced attention are both enabled, the Memory Efficient Attention - is used. - """ - self.unet.set_use_memory_efficient_attention_xformers(True) - - def disable_xformers_memory_efficient_attention(self): - r""" - Disable memory efficient attention as implemented in xformers. - """ - self.unet.set_use_memory_efficient_attention_xformers(False) - def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"): r""" Enable sliced attention computation. diff --git a/pipelines/util.py b/pipelines/util.py new file mode 100644 index 0000000..661dbee --- /dev/null +++ b/pipelines/util.py @@ -0,0 +1,9 @@ +import torch + + +def set_use_memory_efficient_attention_xformers(module: torch.nn.Module, valid: bool) -> None: + if hasattr(module, "set_use_memory_efficient_attention_xformers"): + module.set_use_memory_efficient_attention_xformers(valid) + + for child in module.children(): + set_use_memory_efficient_attention_xformers(child, valid) -- cgit v1.2.3-70-g09d2