diff --git a/toolkit/data_transfer_object/data_loader.py b/toolkit/data_transfer_object/data_loader.py index 7ddc3ce..1111e24 100644 --- a/toolkit/data_transfer_object/data_loader.py +++ b/toolkit/data_transfer_object/data_loader.py @@ -11,6 +11,7 @@ from toolkit import image_utils from toolkit.basic import get_quick_signature_string from toolkit.dto import DTO from toolkit.dataloader_mixins import ( + _latent_from_uint8, CaptionProcessingDTOMixin, ImageProcessingDTOMixin, LatentCachingFileItemDTOMixin, @@ -212,13 +213,34 @@ class FileItemDTO( self.cleanup_mask() self.cleanup_unconditional() + @property + def tensor(self) -> Union[torch.Tensor, None]: + t = self.__dict__.get('_tensor') + if t is None: + cached = getattr(self, '_cached_tensor_uint8', None) + if cached is not None: + # cached pixels stay uint8 until first read, so workers ship 1 byte/px over shm; + # fp16 on read (consumers cast up themselves) + t = _latent_from_uint8(cached, dtype=torch.float16) + self._tensor = t + return t + + @tensor.setter + def tensor(self, value): + self._tensor = value + + @tensor.deleter + def tensor(self): + self._tensor = None + class DataLoaderBatchDTO: def __init__(self, **kwargs): try: self.file_items: List["FileItemDTO"] = kwargs.get("file_items", None) is_latents_cached = self.file_items[0].is_latent_cached - self.tensor: Union[torch.Tensor, None] = None + self._tensor: Union[torch.Tensor, None] = None + self._tensor_uint8: Union[torch.Tensor, None] = None self.latents: Union[torch.Tensor, None] = None self.control_tensor: Union[torch.Tensor, None] = None self.control_tensor_list: Union[List[List[torch.Tensor]], None] = None @@ -265,9 +287,11 @@ class DataLoaderBatchDTO: ): # 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( - [x.tensor.unsqueeze(0) for x in self.file_items] - ) + if all(x.__dict__.get('_tensor') is None and getattr(x, '_cached_tensor_uint8', None) is not None for x in self.file_items): + # cached pixels cross the worker boundary as uint8; `tensor` expands them on first read + self._tensor_uint8 = torch.cat([x._cached_tensor_uint8.unsqueeze(0) for x in self.file_items]) + else: + self.tensor = torch.cat([x.tensor.unsqueeze(0) for x in self.file_items]) # if we have encoded latents, we concatenate them self.latents: Union[torch.Tensor, None] = None if is_latents_cached: @@ -560,6 +584,23 @@ class DataLoaderBatchDTO: def latents(self): self._latents = None + @property + def tensor(self) -> Union[torch.Tensor, None]: + if self._tensor is None and self._tensor_uint8 is not None: + self._tensor = _latent_from_uint8(self._tensor_uint8, dtype=torch.float16) + self._tensor_uint8 = None + return self._tensor + + @tensor.setter + def tensor(self, value): + self._tensor = value + self._tensor_uint8 = None + + @tensor.deleter + def tensor(self): + self._tensor = None + self._tensor_uint8 = None + @property def audio_latents(self) -> Union[torch.Tensor, None]: # cached audio rides inside the latents DTO diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index e10b455..770e691 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -2000,9 +2000,7 @@ class LatentCachingFileItemDTOMixin: if 'waveform' in state_dict: self._cached_waveform_int16 = state_dict['waveform'] self._cached_waveform_sample_rate = int(state_dict['waveform_sample_rate'].item()) - if self._cached_tensor_uint8 is not None and getattr(self, 'tensor', None) is None: - # rebuild the pixel tensor as it would be if loaded without caching - self.tensor = _latent_from_uint8(self._cached_tensor_uint8) + # the pixel tensor is rebuilt from _cached_tensor_uint8 by FileItemDTO.tensor on first read if self._cached_waveform_int16 is not None and self.audio_data is None: # rebuild the audio waveform as it would be if loaded without caching waveform = _waveform_from_int16(self._cached_waveform_int16)