github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add dataset flag to cache the raw tensors

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-04 15:32:48 -06:00
parent 8a0bcf1ffe
commit 00a93e3830
3 changed files with 65 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
toolkit/data_transfer_object/data_loader.py
View File

@@ -224,7 +224,11 @@ class DataLoaderBatchDTO:
self.num_frames: int = self.file_items[0].num_frames
if not is_latents_cached or self.file_items[0].dataset_config.load_image_when_caching_latents:
if (
not is_latents_cached
or self.file_items[0].dataset_config.load_image_when_caching_latents
or self.file_items[0].dataset_config.cache_tensors_to_disk
):
# only return a tensor if latents are not cached, or if we are explicitly
# loading the raw image alongside the cached latents
self.tensor: torch.Tensor = torch.cat(