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

Added a flag to keep loading the image when latents are cached. Useful for DFE and other methods that target pixelspace losses.

Browse Source
This commit is contained in:
Jaret Burkett
2026-06-15 05:31:48 -06:00
parent faa770fc79
commit c730d64478
3 changed files with 18 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
toolkit/dataloader_mixins.py
View File

@@ -754,17 +754,20 @@ class ImageProcessingDTOMixin:
# if we are caching latents, just do that
if self.is_latent_cached:
self.get_latent()
if self.has_control_image:
self.load_control_image()
if self.has_inpaint_image:
self.load_inpaint_image()
if self.has_clip_image:
self.load_clip_image()
if self.has_mask_image:
self.load_mask_image()
if self.has_unconditional:
self.load_unconditional_image()
return
# if load_image_when_caching_latents is set, we still need the raw image
# tensor in addition to the cached latent, so fall through to load it below
if not self.dataset_config.load_image_when_caching_latents:
if self.has_control_image:
self.load_control_image()
if self.has_inpaint_image:
self.load_inpaint_image()
if self.has_clip_image:
self.load_clip_image()
if self.has_mask_image:
self.load_mask_image()
if self.has_unconditional:
self.load_unconditional_image()
return
if self.is_audio_model:
self.load_and_process_audio()
return