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

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)