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

Migrate to a new DTO for latents to carry more information that a normal tensor such as audio.

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-30 10:30:21 -06:00
parent 2a69c1e7de
commit 764b5064fb
8 changed files with 369 additions and 215 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
toolkit/dataloader_mixins.py
View File

@@ -23,6 +23,7 @@ from toolkit.basic import flush, value_map
from toolkit.buckets import get_bucket_for_image_size, get_resolution
from toolkit.config_modules import ControlTypes
from toolkit.control_generator import ControlGenerator
from toolkit.dto import DTO, DISK_PREFIX
from toolkit.metadata import get_meta_for_safetensors
from toolkit.models.pixtral_vision import PixtralVisionImagePreprocessorCompatible
from toolkit.prompt_utils import inject_trigger_into_prompt
@@ -1772,14 +1773,26 @@ def _waveform_from_int16(waveform: torch.Tensor, dtype: torch.dtype = torch.floa
return (waveform.to(torch.float32) / 32767.0).to(dtype)
def _dto_extras_from_state_dict(state_dict) -> dict:
"""Extra latent streams in a cache file: legacy named keys written by
older versions plus the generic dto.<name> keys new caches write."""
extras = {}
if 'audio_latent' in state_dict:
extras['audio'] = state_dict['audio_latent']
for k, v in state_dict.items():
if k.startswith(DISK_PREFIX):
extras[k[len(DISK_PREFIX):]] = v
return extras
class LatentCachingFileItemDTOMixin:
def __init__(self, *args, **kwargs):
# if we have super, call it
if hasattr(super(), '__init__'):
super().__init__(*args, **kwargs)
# a plain tensor, or a DTO carrying extra streams (audio rows, ...)
self._encoded_latent: Union[torch.Tensor, None] = None
self._cached_first_frame_latent: Union[torch.Tensor, None] = None
self._cached_audio_latent: Union[torch.Tensor, None] = None
self._cached_tensor_uint8: Union[torch.Tensor, None] = None
self._cached_waveform_int16: Union[torch.Tensor, None] = None
self._cached_waveform_sample_rate: Union[int, None] = None
@@ -1862,17 +1875,14 @@ class LatentCachingFileItemDTOMixin:
# we are caching on disk, don't save in memory
self._encoded_latent = None
self._cached_first_frame_latent = None
self._cached_audio_latent = None
self._cached_tensor_uint8 = None
self._cached_waveform_int16 = None
self._cached_waveform_sample_rate = None
else:
# move it back to cpu
# move it back to cpu (a DTO carries its extras along)
self._encoded_latent = self._encoded_latent.to('cpu')
if self._cached_first_frame_latent is not None:
self._cached_first_frame_latent = self._cached_first_frame_latent.to('cpu')
if self._cached_audio_latent is not None:
self._cached_audio_latent = self._cached_audio_latent.to('cpu')
def get_latent(self, device=None):
if not self.is_latent_cached:
@@ -1892,8 +1902,9 @@ class LatentCachingFileItemDTOMixin:
self._cached_first_frame_latent = state_dict['first_frame_latent']
if self._cached_first_frame_latent.dtype == torch.uint8:
self._cached_first_frame_latent = _latent_from_uint8(self._cached_first_frame_latent)
if 'audio_latent' in state_dict:
self._cached_audio_latent = state_dict['audio_latent']
extras = _dto_extras_from_state_dict(state_dict)
if extras:
self._encoded_latent = DTO(self._encoded_latent, **extras)
if 'num_frames' in state_dict:
self.num_frames = int(state_dict['num_frames'].item())
if 'tensor' in state_dict:
@@ -2008,14 +2019,15 @@ class LatentCachingMixin:
if cached_latent.dtype == torch.uint8:
# pixel-space latents cached as uint8
cached_latent = _latent_from_uint8(cached_latent)
extras = _dto_extras_from_state_dict(state_dict)
if extras:
cached_latent = DTO(cached_latent, **extras)
file_item._encoded_latent = cached_latent.to('cpu', dtype=self.sd.torch_dtype)
if 'first_frame_latent' in state_dict:
cached_first_frame = state_dict['first_frame_latent']
if cached_first_frame.dtype == torch.uint8:
cached_first_frame = _latent_from_uint8(cached_first_frame)
file_item._cached_first_frame_latent = cached_first_frame.to('cpu', dtype=self.sd.torch_dtype)
if 'audio_latent' in state_dict:
file_item._cached_audio_latent = state_dict['audio_latent'].to('cpu', dtype=self.sd.torch_dtype)
if 'tensor' in state_dict:
file_item._cached_tensor_uint8 = state_dict['tensor']
if 'waveform' in state_dict:
@@ -2051,12 +2063,19 @@ class LatentCachingMixin:
file_item._cached_waveform_sample_rate = sample_rate
try:
imgs = file_item.tensor.unsqueeze(0).to(device, dtype=dtype)
latent = self.sd.encode_images(imgs).squeeze(0)
latent = self.sd.encode_images(imgs)
# a model can return a DTO carrying extra streams alongside the latent
latent = latent.map(lambda t: t.squeeze(0)) if isinstance(latent, DTO) else latent.squeeze(0)
if to_disk:
main_latent = latent.tensor if isinstance(latent, DTO) else latent
if cache_uint8:
state_dict['latent'] = _latent_to_uint8(latent).cpu()
state_dict['latent'] = _latent_to_uint8(main_latent).cpu()
else:
state_dict['latent'] = latent.clone().detach().cpu()
state_dict['latent'] = main_latent.clone().detach().cpu()
if isinstance(latent, DTO):
for k, v in latent.extras.items():
if torch.is_tensor(v):
state_dict[f'{DISK_PREFIX}{k}'] = v.clone().detach().cpu()
except Exception as e:
print_acc(f"Error processing image: {file_item.path}")
print_acc(f"Error: {str(e)}")
@@ -2095,12 +2114,12 @@ class LatentCachingMixin:
save_file(state_dict, latent_path, metadata=meta)
if to_memory:
# keep it in memory
# keep it in memory; audio rides inside the latent DTO
if audio_latent is not None:
latent = DTO(latent, audio=audio_latent)
file_item._encoded_latent = latent.to('cpu', dtype=self.sd.torch_dtype)
if first_frame_latent is not None:
file_item._cached_first_frame_latent = first_frame_latent.to('cpu', dtype=self.sd.torch_dtype)
if audio_latent is not None:
file_item._cached_audio_latent = audio_latent.to('cpu', dtype=self.sd.torch_dtype)
del imgs
del latent