From 6785331e22d6a488fbf5905fab56d7fec867e038 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Sun, 2 Oct 2022 22:59:01 +0300
Subject: [PATCH] keep textual inversion dataset latents in CPU memory to save
 a bit of VRAM

---
 modules/textual_inversion/dataset.py           | 2 ++
 modules/textual_inversion/textual_inversion.py | 3 +++
 modules/ui.py                                  | 4 ++--
 3 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/modules/textual_inversion/dataset.py b/modules/textual_inversion/dataset.py
index 7e134a08..e8394ff6 100644
--- a/modules/textual_inversion/dataset.py
+++ b/modules/textual_inversion/dataset.py
@@ -8,6 +8,7 @@ from torchvision import transforms
 
 import random
 import tqdm
+from modules import devices
 
 
 class PersonalizedBase(Dataset):
@@ -47,6 +48,7 @@ class PersonalizedBase(Dataset):
             torchdata = torch.moveaxis(torchdata, 2, 0)
 
             init_latent = model.get_first_stage_encoding(model.encode_first_stage(torchdata.unsqueeze(dim=0))).squeeze()
+            init_latent = init_latent.to(devices.cpu)
 
             self.dataset.append((init_latent, filename_tokens))
 
diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index d4e250d8..8686f534 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -212,7 +212,10 @@ def train_embedding(embedding_name, learn_rate, data_root, log_directory, steps,
 
         with torch.autocast("cuda"):
             c = cond_model([text])
+
+            x = x.to(devices.device)
             loss = shared.sd_model(x.unsqueeze(0), c)[0]
+            del x
 
             losses[embedding.step % losses.shape[0]] = loss.item()
 
diff --git a/modules/ui.py b/modules/ui.py
index e7bde53b..d9d02ece 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -1002,8 +1002,8 @@ def create_ui(wrap_gradio_gpu_call):
                     log_directory = gr.Textbox(label='Log directory', placeholder="Path to directory where to write outputs", value="textual_inversion")
                     template_file = gr.Textbox(label='Prompt template file', value=os.path.join(script_path, "textual_inversion_templates", "style_filewords.txt"))
                     steps = gr.Number(label='Max steps', value=100000, precision=0)
-                    create_image_every = gr.Number(label='Save an image to log directory every N steps, 0 to disable', value=1000, precision=0)
-                    save_embedding_every = gr.Number(label='Save a copy of embedding to log directory every N steps, 0 to disable', value=1000, precision=0)
+                    create_image_every = gr.Number(label='Save an image to log directory every N steps, 0 to disable', value=500, precision=0)
+                    save_embedding_every = gr.Number(label='Save a copy of embedding to log directory every N steps, 0 to disable', value=500, precision=0)
 
                     with gr.Row():
                         with gr.Column(scale=2):