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

don't hold latent on memory for finetuning dataset

Browse Source
This commit is contained in:
Kohya S
2023-07-10 08:46:15 +09:00
parent 5c80117fbd
commit b6e328ea8f
1 changed files with 4 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
library/train_util.py
View File

@@ -91,6 +91,7 @@ IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".webp", ".bmp", ".PNG", ".JPG", ".
try: try:
import pillow_avif import pillow_avif
IMAGE_EXTENSIONS.extend([".avif", ".AVIF"]) IMAGE_EXTENSIONS.extend([".avif", ".AVIF"])
except: except:
pass pass
@@ -853,16 +854,11 @@ class BaseDataset(torch.utils.data.Dataset):
# split by resolution # split by resolution
batches = [] batches = []
batch = [] batch = []
for info in image_infos: print("checking cache validity...")
for info in tqdm(image_infos):
subset = self.image_to_subset[info.image_key] subset = self.image_to_subset[info.image_key]
if info.latents_npz is not None: if info.latents_npz is not None: # fine tuning dataset
info.latents, info.latents_original_size, info.latents_crop_left_top = self.load_latents_from_npz(info, False)
info.latents = torch.FloatTensor(info.latents)
info.latents_flipped, _, _ = self.load_latents_from_npz(info, True) # might be None
if info.latents_flipped is not None:
info.latents_flipped = torch.FloatTensor(info.latents_flipped)
continue continue
# check disk cache exists and size of latents # check disk cache exists and size of latents