diff options
| author | AUTOMATIC <16777216c@gmail.com> | 2022-10-02 19:59:01 +0000 |
|---|---|---|
| committer | AUTOMATIC <16777216c@gmail.com> | 2022-10-02 19:59:01 +0000 |
| commit | 6785331e22d6a488fbf5905fab56d7fec867e038 (patch) | |
| tree | 111f08d07e3dfc82a02857155cccdeefe7afdb70 /modules/textual_inversion/dataset.py | |
| parent | c7543d4940da672d970124ae8f2fec9de7bdc1da (diff) | |
| download | stable-diffusion-webui-gfx803-6785331e22d6a488fbf5905fab56d7fec867e038.tar.gz stable-diffusion-webui-gfx803-6785331e22d6a488fbf5905fab56d7fec867e038.tar.bz2 stable-diffusion-webui-gfx803-6785331e22d6a488fbf5905fab56d7fec867e038.zip | |
keep textual inversion dataset latents in CPU memory to save a bit of VRAM
Diffstat (limited to 'modules/textual_inversion/dataset.py')
| -rw-r--r-- | modules/textual_inversion/dataset.py | 2 |
1 files changed, 2 insertions, 0 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))
|