VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-09 06:45:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

cache latents to disk in dreambooth method

Browse Source
This commit is contained in:
Kohya S
2023-04-12 23:10:39 +09:00
parent 5050971ac6
commit 2e9f7b5f91
6 changed files with 67 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
train_network.py
View File

@@ -172,12 +172,14 @@ def train(args):
vae.requires_grad_(False)
vae.eval()
with torch.no_grad():
train_dataset_group.cache_latents(vae, args.vae_batch_size)
train_dataset_group.cache_latents(vae, args.vae_batch_size, args.cache_latents_to_disk, accelerator.is_main_process)
vae.to("cpu")
if torch.cuda.is_available():
torch.cuda.empty_cache()
gc.collect()
accelerator.wait_for_everyone()
# prepare network
import sys