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

3
toolkit/config_modules.py
View File

@@ -994,6 +994,9 @@ class DatasetConfig:
self.cache_latents: bool = kwargs.get('cache_latents', False)
# cache latents to disk will store them on disk. If both are true, it will save to disk, but keep in memory
self.cache_latents_to_disk: bool = kwargs.get('cache_latents_to_disk', False)
# cache tensors to disk. Useful for saving video files tensors to the disk so we have the clean pixelspace versions of video and audio
self.cache_tensors_to_disk: bool = kwargs.get('cache_tensors_to_disk', False)
self.cache_clip_vision_to_disk: bool = kwargs.get('cache_clip_vision_to_disk', False)
self.cache_text_embeddings: bool = kwargs.get('cache_text_embeddings', False)
self.load_image_when_caching_latents: bool = kwargs.get('load_image_when_caching_latents', False)