1 files changed, 9 insertions, 0 deletions

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)