Buckets:
| 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 | |
| """ | |
| 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 | |
| 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)) | |
| def tensor(self) -> torch.Tensor: | |
| with torch._C.DisableTorchFunctionSubclass(): | |
| return self.as_subclass(torch.Tensor) | |
| 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()) | |
| 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) | |
| 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 | |
| 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) | |
Xet Storage Details
- Size:
- 7.55 kB
- Xet hash:
- 9e36e34c956988a84818cdcea7c458c3f5453a214318bb6d4331ec689c5df367
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.