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

toolkit/config_modules.py

@@ -433,7 +433,6 @@ class TrainConfig:
         self.random_noise_shift = kwargs.get('random_noise_shift', 0.0)
         self.img_multiplier = kwargs.get('img_multiplier', 1.0)
         self.noisy_latent_multiplier = kwargs.get('noisy_latent_multiplier', 1.0)
-        self.latent_multiplier = kwargs.get('latent_multiplier', 1.0)
         self.negative_prompt = kwargs.get('negative_prompt', None)
         self.max_negative_prompts = kwargs.get('max_negative_prompts', 1)
         # multiplier applied to loos on regularization images
@@ -502,7 +501,6 @@ class TrainConfig:
 
         # standardize inputs to the meand std of the model knowledge
         self.standardize_images = kwargs.get('standardize_images', False)
-        self.standardize_latents = kwargs.get('standardize_latents', False)
 
         # if self.train_turbo and not self.noise_scheduler.startswith("euler"):
         #     raise ValueError(f"train_turbo is only supported with euler and wuler_a noise schedulers")

toolkit/data_transfer_object/data_loader.py

@@ -9,6 +9,7 @@ import av
             
 from toolkit import image_utils
 from toolkit.basic import get_quick_signature_string
+from toolkit.dto import DTO
 from toolkit.dataloader_mixins import (
     CaptionProcessingDTOMixin,
     ImageProcessingDTOMixin,
@@ -239,7 +240,6 @@ class DataLoaderBatchDTO:
             )
             self.audio_tensor: Union[torch.Tensor, None] = None
             self.first_frame_latents: Union[torch.Tensor, None] = None
-            self.audio_latents: Union[torch.Tensor, None] = None
             # control-video reference paths (encoded + disk-cached lazily by
             # models with supports_video_control_images)
             self.control_video_paths_list: Union[List, None] = None
@@ -249,30 +249,8 @@ class DataLoaderBatchDTO:
                     for x in self.file_items
                 ]
 
-            # just for holding noise and preds during training
-            self.audio_target: Union[torch.Tensor, None] = None
-            self.audio_pred: Union[torch.Tensor, None] = None
-            # the noise drawn for the audio stream on the primary (grad enabled)
-            # prediction. Secondary passes (cfg / guidance loss / prior preds)
-            # reuse it so their noisy audio matches the stored audio_target.
-            self.audio_noise: Union[torch.Tensor, None] = None
-            # audio predictions from the non primary passes. Kept separate so
-            # they cannot stomp the primary pred we backprop through.
-            self.audio_pred_uncond: Union[torch.Tensor, None] = None
-            self.audio_pred_prior: Union[torch.Tensor, None] = None
-            self.audio_pred_preservation: Union[torch.Tensor, None] = None
-            # which of the above the current secondary pass writes to. None (the
-            # default) means no secondary pass is in flight: any grad-enabled
-            # prediction is a primary one and writes audio_pred (and the
-            # noisy/sigma bookkeeping) directly. The trainer sets this around
-            # its prior / guidance-unconditional / preservation passes.
-            self.audio_pred_slot: Union[str, None] = None
             # set by the trainer around the D-OPSD teacher pass
             self.dopsd_teacher_pass: bool = False
-            # noisy audio rows and audio sigma of the primary pass, used to
-            # rebuild the clean audio estimate for perceptual losses
-            self.audio_noisy: Union[torch.Tensor, None] = None
-            self.audio_sigma: Union[torch.Tensor, None] = None
 
             self.num_frames: int = self.file_items[0].num_frames
 
@@ -289,10 +267,9 @@ class DataLoaderBatchDTO:
             # if we have encoded latents, we concatenate them
             self.latents: Union[torch.Tensor, None] = None
             if is_latents_cached:
-                # this get_latent call with trigger loading all cached items from the disk
-                self.latents = torch.cat(
-                    [x.get_latent().unsqueeze(0) for x in self.file_items]
-                )
+                # this get_latent call with trigger loading all cached items from the disk.
+                # DTO.stack keeps any extra streams (audio, ...) riding on the batch latents
+                self.latents = DTO.stack([x.get_latent() for x in self.file_items])
                 if any(
                     [x._cached_first_frame_latent is not None for x in self.file_items]
                 ):
@@ -310,22 +287,6 @@ class DataLoaderBatchDTO:
                             for x in self.file_items
                         ]
                     )
-                if any([x._cached_audio_latent is not None for x in self.file_items]):
-                    # find one to use as a base; item 0 may not have one
-                    base_audio_latent = None
-                    for x in self.file_items:
-                        if x._cached_audio_latent is not None:
-                            base_audio_latent = x._cached_audio_latent
-                            break
-                    self.audio_latents = torch.cat(
-                        [
-                            x._cached_audio_latent.unsqueeze(0)
-                            if x._cached_audio_latent is not None
-                            else torch.zeros_like(base_audio_latent).unsqueeze(0)
-                            for x in self.file_items
-                        ]
-                    )
-
             self.prompt_embeds: Union[PromptEmbeds, None] = None
             # diff output preservation embeds (trigger word replaced with class)
             self.dop_prompt_embeds: Union[PromptEmbeds, None] = None
@@ -576,14 +537,31 @@ class DataLoaderBatchDTO:
     ):
         return [x.caption_short for x in self.file_items]
 
-    def set_secondary_audio_pred(self, pred):
-        """Route an audio prediction from a non primary pass (prior,
-        unconditional/guidance, preservation) to its own slot so it cannot
-        stomp the primary prediction the loss backprops through. Passes that
-        did not declare a slot (e.g. a trainer's extra no_grad prediction)
-        are simply not stored."""
-        if self.audio_pred_slot is not None:
-            setattr(self, self.audio_pred_slot, pred)
+    @property
+    def latents(self) -> Union[torch.Tensor, None]:
+        return self._latents
+
+    @latents.setter
+    def latents(self, value):
+        # math on a DTO returns a plain tensor, so trainer code that rescales
+        # or re-noises the latents would silently drop the extra streams
+        # (audio, ...) riding on them. Re-attach them here so assignments stay
+        # plain tensor code everywhere else.
+        prev = getattr(self, '_latents', None)
+        if isinstance(prev, DTO) and value is not None and not isinstance(value, DTO):
+            value = DTO(value, **prev.extras)
+        self._latents = value
+
+    @latents.deleter
+    def latents(self):
+        self._latents = None
+
+    @property
+    def audio_latents(self) -> Union[torch.Tensor, None]:
+        # cached audio rides inside the latents DTO
+        if isinstance(self.latents, DTO):
+            return self.latents.get('audio')
+        return None
 
     def cleanup(self):
         del self.latents
@@ -591,16 +569,7 @@ class DataLoaderBatchDTO:
         del self.control_tensor
         del self.audio_tensor
         del self.audio_data
-        del self.audio_target
-        del self.audio_pred
-        del self.audio_noise
-        del self.audio_pred_uncond
-        del self.audio_pred_prior
-        del self.audio_pred_preservation
-        del self.audio_noisy
-        del self.audio_sigma
         del self.first_frame_latents
-        del self.audio_latents
         for file_item in self.file_items:
             file_item.cleanup()
 

toolkit/dataloader_mixins.py

@@ -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

toolkit/dto.py

@@ -0,0 +1,208 @@
+import torch
+
+# extra tensors ride in the latent cache safetensors under this key prefix
+DISK_PREFIX = "dto."
+
+
+def _unwrap(value):
+    """Recursively convert any DTO in value back to a plain torch.Tensor.
+    Internal: everywhere else, use `dto.tensor`."""
+    if isinstance(value, DTO):
+        return value.tensor
+    if isinstance(value, (list, tuple)):
+        return type(value)(_unwrap(v) for v in value)
+    return value
+
+
+def _rebuild_dto(tensor, extras):
+    return DTO(tensor, **extras)
+
+
+class DTO(torch.Tensor):
+    """A torch.Tensor that carries named side-channel data (audio rows, video
+    tokens, extra targets, ...) through code that only knows about the main
+    tensor.
+
+    Backwards compatible by design: a DTO *is* the main tensor, so every
+    existing shape check, math op, and indexing keeps working. Any torch op
+    returns a plain tensor — extras never leak through math — and only the
+    explicit carriers (``to``, ``clone``, ``detach``, ``cpu``, ``cuda``,
+    ``pin_memory``, ``map``, ``cat``) keep the extras attached.
+
+        latent = DTO(video_latent, audio=audio_rows, num_frames=77)
+        latent.audio          # extra lookup, AttributeError if missing
+        latent.get("audio")   # None if missing
+        latent.tensor         # plain tensor view (shares storage)
+        latent * 2            # plain tensor, extras dropped
+        latent.to("cuda")     # DTO, tensor extras moved too
+    """
+
+    @staticmethod
+    def __new__(cls, tensor: torch.Tensor, **extras):
+        if isinstance(tensor, DTO):
+            extras = {**tensor.extras, **extras}
+            tensor = tensor.tensor
+        obj = tensor.as_subclass(cls)
+        obj._dto_extras = dict(extras)
+        return obj
+
+    @classmethod
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
+        # run every torch op as if on plain tensors so extras never
+        # accidentally propagate through math with stale values
+        if kwargs is None:
+            kwargs = {}
+        with torch._C.DisableTorchFunctionSubclass():
+            return _unwrap(func(*args, **kwargs))
+
+    @property
+    def tensor(self) -> torch.Tensor:
+        with torch._C.DisableTorchFunctionSubclass():
+            return self.as_subclass(torch.Tensor)
+
+    @property
+    def extras(self) -> dict:
+        return self._dto_extras
+
+    def get(self, key, default=None):
+        return self._dto_extras.get(key, default)
+
+    def set(self, key, value):
+        self._dto_extras[key] = value
+        return self
+
+    def __getattr__(self, name):
+        if name == "_dto_extras":
+            raise AttributeError(name)
+        try:
+            return self._dto_extras[name]
+        except KeyError:
+            raise AttributeError(
+                f"DTO has no extra '{name}'; extras: {list(self._dto_extras.keys())}"
+            )
+
+    def __repr__(self):
+        return f"DTO(extras={list(self._dto_extras.keys())}, tensor={self.tensor!r})"
+
+    def __reduce_ex__(self, protocol):
+        return (_rebuild_dto, (self.tensor, self._dto_extras))
+
+    def map(self, fn):
+        """Apply fn to the main tensor and every tensor extra, keep the rest."""
+        return DTO(
+            fn(self.tensor),
+            **{
+                k: fn(v) if torch.is_tensor(v) else v
+                for k, v in self._dto_extras.items()
+            },
+        )
+
+    def _carry(self, base, fn_tensor):
+        return DTO(
+            base,
+            **{
+                k: fn_tensor(v) if torch.is_tensor(v) else v
+                for k, v in self._dto_extras.items()
+            },
+        )
+
+    def to(self, *args, **kwargs):
+        device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
+
+        def move(t):
+            # dtype casts only follow onto floating extras; int extras
+            # (frame counts, indices) keep their dtype on device moves
+            d = dtype if dtype is not None and t.is_floating_point() else None
+            return t.to(device=device, dtype=d, non_blocking=non_blocking)
+
+        return self._carry(self.tensor.to(*args, **kwargs), move)
+
+    def cpu(self):
+        return self.to("cpu")
+
+    def cuda(self, device=None):
+        return self.to(device if device is not None else "cuda")
+
+    def clone(self):
+        return self.map(lambda t: t.clone())
+
+    def detach(self):
+        return self.map(lambda t: t.detach())
+
+    def pin_memory(self):
+        return self.map(lambda t: t.pin_memory())
+
+    @classmethod
+    def cat(cls, items, dim=0):
+        """Batch-collate: cat main tensors along dim, tensor extras shared by
+        every item along dim 0. Non-tensor extras keep a single value when
+        identical everywhere, else become a list."""
+        base = torch.cat([_unwrap(x) for x in items], dim=dim)
+        dtos = [x for x in items if isinstance(x, cls)]
+        if len(dtos) != len(items):
+            return base if not dtos else cls(base, **dtos[0].extras)
+        keys = set(dtos[0].extras.keys())
+        for d in dtos[1:]:
+            keys &= set(d.extras.keys())
+        extras = {}
+        for k in keys:
+            vals = [d.extras[k] for d in dtos]
+            if all(torch.is_tensor(v) for v in vals):
+                extras[k] = torch.cat(vals, dim=0)
+            elif all(v == vals[0] for v in vals[1:]) if len(vals) > 1 else True:
+                extras[k] = vals[0]
+            else:
+                extras[k] = vals
+        return cls(base, **extras)
+
+    @classmethod
+    def stack(cls, items):
+        """Collate per-item latents into a batch: unsqueeze(0) + cat. A tensor
+        extra missing on some items is zero-filled there (a missing stream is
+        silence). Returns a plain tensor when no item carries extras."""
+        base = torch.cat([_unwrap(x).unsqueeze(0) for x in items], dim=0)
+        keys = []
+        for x in items:
+            if isinstance(x, cls):
+                keys.extend(k for k in x.extras if k not in keys)
+        if not keys:
+            return base
+        extras = {}
+        for k in keys:
+            vals = [x.get(k) if isinstance(x, cls) else None for x in items]
+            present = [v for v in vals if v is not None]
+            if all(torch.is_tensor(v) for v in present):
+                extras[k] = torch.cat(
+                    [
+                        (v if v is not None else torch.zeros_like(present[0])).unsqueeze(0)
+                        for v in vals
+                    ],
+                    dim=0,
+                )
+            elif all(v == present[0] for v in present[1:]):
+                extras[k] = present[0]
+            else:
+                extras[k] = vals
+        return cls(base, **extras)
+
+    def to_state_dict(self, key="latent") -> dict:
+        """Flatten for safetensors: main tensor under ``key``, tensor extras
+        under ``dto.<name>``. Non-tensor extras are not persisted."""
+        state_dict = {key: self.tensor.contiguous()}
+        for k, v in self._dto_extras.items():
+            if torch.is_tensor(v):
+                state_dict[f"{DISK_PREFIX}{k}"] = v.contiguous()
+        return state_dict
+
+    @staticmethod
+    def from_state_dict(state_dict: dict, key="latent"):
+        """Inverse of ``to_state_dict``. Returns a plain tensor when the file
+        holds no dto extras, so legacy caches load unchanged."""
+        extras = {
+            k[len(DISK_PREFIX):]: v
+            for k, v in state_dict.items()
+            if k.startswith(DISK_PREFIX)
+        }
+        if not extras:
+            return state_dict[key]
+        return DTO(state_dict[key], **extras)