From 2825ba2f2030b2fd3e841aad416a4fd28d67615a Mon Sep 17 00:00:00 2001 From: Volpeon Date: Sun, 4 Dec 2022 09:24:37 +0100 Subject: Update --- textual_inversion.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) (limited to 'textual_inversion.py') 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() -- cgit v1.2.3-54-g00ecf