Update

1 files changed, 3 insertions, 1 deletions

diff --git a/textual_inversion.py b/textual_inversion.py
index cd2d22b..1a5a8d0 100644
--- a/textual_inversion.py
+++ b/textual_inversion.py
@@ -23,6 +23,7 @@ from transformers import CLIPTextModel, CLIPTokenizer
 from slugify import slugify
 
 from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
+from pipelines.util import set_use_memory_efficient_attention_xformers
 from data.csv import CSVDataModule
 from training.optimization import get_one_cycle_schedule
 from models.clip.prompt import PromptProcessor
@@ -546,7 +547,8 @@ def main():
         args.pretrained_model_name_or_path, subfolder='scheduler')
 
     vae.enable_slicing()
-    unet.set_use_memory_efficient_attention_xformers(True)
+    set_use_memory_efficient_attention_xformers(unet, True)
+    set_use_memory_efficient_attention_xformers(vae, True)
 
     if args.gradient_checkpointing:
         text_encoder.gradient_checkpointing_enable()