From a9a04547e92a51583999ea2e7da8792da80e0d7a Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Mon, 3 Aug 2026 15:54:30 -0600 Subject: [PATCH] Dont move encoders on and off device when caching until the first instance of needing to process a cache item. --- toolkit/dataloader_mixins.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 02c7afa1..19ef5d07 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -1811,8 +1811,9 @@ class LatentCachingMixin: print_acc(" - Saving latents to disk") if to_memory: print_acc(" - Keeping latents in memory") - # move sd items to cpu except for vae - self.sd.set_device_state_preset('cache_latents') + # move sd items to cpu except for vae. Only done on the first item that + # actually needs encoding so fully cached datasets don't shuffle models around + did_move = False # prep (video decode, frame extraction, audio load, disk reads) is done by a # thread pool so the next items are ready while the current one is encoding. @@ -1852,6 +1853,9 @@ class LatentCachingMixin: next_item = next(file_iter, None) if next_item is not None: pending.append(executor.submit(_prep, next_item)) + if needs_encode and not did_move: + self.sd.set_device_state_preset('cache_latents') + did_move = True self._cache_one_latent(file_item, latent_path, cached_state_dict, needs_encode, to_disk, to_memory) file_item.is_latent_cached = True i += 1 @@ -1861,7 +1865,8 @@ class LatentCachingMixin: pbar.close() # restore device state - self.sd.restore_device_state() + if did_move: + self.sd.restore_device_state() def _cache_one_latent( self: 'AiToolkitDataset',