| from __future__ import annotations |
|
|
| from bisect import bisect_right |
| from pathlib import Path |
| from typing import Iterator |
|
|
| import torch |
| from torch import Tensor |
| from torch.utils.data import Dataset, Sampler |
|
|
|
|
| def _load_shard(path: Path) -> dict[str, object]: |
| try: |
| value = torch.load(path, map_location="cpu", weights_only=True, mmap=True) |
| except TypeError: |
| |
| value = torch.load(path, map_location="cpu", weights_only=True) |
| if not isinstance(value, dict): |
| raise ValueError(f"{path}: expected a dictionary of tensors") |
| return value |
|
|
|
|
| class LatentTextShardDataset(Dataset[dict[str, Tensor]]): |
| """Lazy reader for tensor shards, with one latent resolution per shard. |
| |
| Required keys: ``latents`` [N,C,H,W], ``text_embeddings`` [N,L,D], and |
| ``text_mask`` [N,L]. Files are loaded with PyTorch mmap when available. |
| """ |
|
|
| def __init__(self, data_dir: str | Path, max_text_tokens: int = 96): |
| self.paths = sorted(Path(data_dir).glob("*.pt")) |
| if not self.paths: |
| raise FileNotFoundError(f"no .pt shard files found in {data_dir}") |
| self.max_text_tokens = max_text_tokens |
| self.metadata: list[dict[str, int]] = [] |
| self.cumulative_ends: list[int] = [] |
| self._cache: dict[int, dict[str, object]] = {} |
| self._cache_limit = 4 |
| total = 0 |
| latent_channels: int | None = None |
| text_dim: int | None = None |
| repa_dim: int | None = None |
| has_teacher_features: bool | None = None |
| for shard_index, path in enumerate(self.paths): |
| shard = _load_shard(path) |
| latents = shard.get("latents") |
| text = shard.get("text_embeddings") |
| mask = shard.get("text_mask") |
| if not all(isinstance(x, Tensor) for x in (latents, text, mask)): |
| raise ValueError(f"{path}: required tensor keys are latents, text_embeddings, text_mask") |
| assert isinstance(latents, Tensor) and isinstance(text, Tensor) and isinstance(mask, Tensor) |
| if latents.ndim != 4 or text.ndim != 3 or mask.ndim != 2: |
| raise ValueError(f"{path}: expected [N,C,H,W], [N,L,D], and [N,L] tensors") |
| count, channels, height, width = latents.shape |
| if count < 1 or text.shape[0] != count or mask.shape != text.shape[:2]: |
| raise ValueError(f"{path}: batch/sample counts do not match") |
| if text.shape[1] > max_text_tokens: |
| raise ValueError(f"{path}: {text.shape[1]} text tokens exceeds limit {max_text_tokens}") |
| if latent_channels is None: |
| latent_channels = channels |
| text_dim = text.shape[2] |
| if channels != latent_channels or text.shape[2] != text_dim: |
| raise ValueError(f"{path}: channel or text embedding dimensions differ from earlier shards") |
| if not latents.is_floating_point() or not text.is_floating_point(): |
| raise ValueError(f"{path}: latents and text_embeddings must be floating point") |
| teacher = shard.get("teacher_features") |
| shard_has_teacher = isinstance(teacher, Tensor) |
| if has_teacher_features is None: |
| has_teacher_features = shard_has_teacher |
| elif shard_has_teacher != has_teacher_features: |
| raise ValueError("either every shard or no shard must contain teacher_features") |
| if shard_has_teacher: |
| assert isinstance(teacher, Tensor) |
| if teacher.ndim != 3 or teacher.shape[0] != count or teacher.shape[1] != height * width: |
| raise ValueError(f"{path}: teacher_features must have shape [N,H*W,feature_dim]") |
| if not teacher.is_floating_point(): |
| raise ValueError(f"{path}: teacher_features must be floating point values") |
| if repa_dim is None: |
| repa_dim = teacher.shape[2] |
| elif teacher.shape[2] != repa_dim: |
| raise ValueError("teacher feature dimensions differ across shards") |
| self.metadata.append( |
| {"start": total, "count": count, "height": height, "width": width, "shard": shard_index} |
| ) |
| total += count |
| self.cumulative_ends.append(total) |
| self.total_samples = total |
| self.latent_channels = int(latent_channels or 0) |
| self.text_dim = int(text_dim or 0) |
| self.has_teacher_features = bool(has_teacher_features) |
| self.repa_dim = int(repa_dim or 0) |
|
|
| def __len__(self) -> int: |
| return self.total_samples |
|
|
| def _locate(self, index: int) -> tuple[int, int]: |
| shard_index = bisect_right(self.cumulative_ends, index) |
| previous_end = self.cumulative_ends[shard_index - 1] if shard_index else 0 |
| return shard_index, index - previous_end |
|
|
| def _get_shard(self, shard_index: int) -> dict[str, object]: |
| if shard_index not in self._cache: |
| self._cache[shard_index] = _load_shard(self.paths[shard_index]) |
| while len(self._cache) > self._cache_limit: |
| del self._cache[next(iter(self._cache))] |
| return self._cache[shard_index] |
|
|
| def __getitem__(self, index: int) -> dict[str, Tensor]: |
| shard_index, row = self._locate(index) |
| shard = self._get_shard(shard_index) |
| result = { |
| "latents": shard["latents"][row].float(), |
| "text_embeddings": shard["text_embeddings"][row].float(), |
| "text_mask": shard["text_mask"][row].bool(), |
| } |
| if self.has_teacher_features: |
| result["teacher_features"] = shard["teacher_features"][row].float() |
| return result |
|
|
| def bucket_shards(self) -> dict[tuple[int, int], list[dict[str, int]]]: |
| buckets: dict[tuple[int, int], list[dict[str, int]]] = {} |
| for meta in self.metadata: |
| buckets.setdefault((meta["height"], meta["width"]), []).append(meta) |
| return buckets |
|
|
|
|
| class ResolutionBucketBatchSampler(Sampler[list[int]]): |
| """Shuffle batches while ensuring every batch has a consistent latent grid.""" |
|
|
| def __init__( |
| self, |
| dataset: LatentTextShardDataset, |
| batch_size: int, |
| *, |
| seed: int = 1234, |
| drop_last: bool = True, |
| rank: int = 0, |
| world_size: int = 1, |
| ): |
| if batch_size < 1: |
| raise ValueError("batch_size must be positive") |
| self.dataset = dataset |
| self.batch_size = batch_size |
| self.seed = seed |
| self.drop_last = drop_last |
| self.rank = rank |
| self.world_size = world_size |
| self.epoch = 0 |
|
|
| def set_epoch(self, epoch: int) -> None: |
| self.epoch = epoch |
|
|
| def _make_batches(self) -> list[list[int]]: |
| generator = torch.Generator().manual_seed(self.seed + self.epoch) |
| buckets = self.dataset.bucket_shards() |
| bucket_items = list(buckets.items()) |
| if bucket_items: |
| order = torch.randperm(len(bucket_items), generator=generator).tolist() |
| bucket_items = [bucket_items[i] for i in order] |
| batches: list[list[int]] = [] |
| for _, shards in bucket_items: |
| pending: list[int] = [] |
| shard_order = torch.randperm(len(shards), generator=generator).tolist() |
| for shard_position in shard_order: |
| meta = shards[shard_position] |
| rows = torch.randperm(meta["count"], generator=generator).tolist() |
| pending.extend(meta["start"] + row for row in rows) |
| while len(pending) >= self.batch_size: |
| batches.append(pending[: self.batch_size]) |
| pending = pending[self.batch_size :] |
| if pending and not self.drop_last: |
| batches.append(pending) |
| if batches: |
| order = torch.randperm(len(batches), generator=generator).tolist() |
| batches = [batches[i] for i in order] |
| if self.world_size > 1: |
| usable = len(batches) - (len(batches) % self.world_size) |
| batches = batches[:usable][self.rank : usable : self.world_size] |
| return batches |
|
|
| def __iter__(self) -> Iterator[list[int]]: |
| yield from self._make_batches() |
|
|
| def __len__(self) -> int: |
| count = 0 |
| for (height, width), shards in self.dataset.bucket_shards().items(): |
| samples = sum(meta["count"] for meta in shards) |
| count += samples // self.batch_size if self.drop_last else (samples + self.batch_size - 1) // self.batch_size |
| if self.world_size > 1: |
| count -= count % self.world_size |
| count //= self.world_size |
| return count |
|
|
|
|
| def collate_latent_text(samples: list[dict[str, Tensor]]) -> dict[str, Tensor]: |
| if not samples: |
| raise ValueError("cannot collate an empty batch") |
| latent_shapes = {tuple(sample["latents"].shape) for sample in samples} |
| if len(latent_shapes) != 1: |
| raise ValueError("all samples in a batch must have the same latent shape") |
| max_tokens = max(sample["text_embeddings"].shape[0] for sample in samples) |
| text_dim = samples[0]["text_embeddings"].shape[-1] |
| batch = len(samples) |
| embeddings = torch.zeros(batch, max_tokens, text_dim, dtype=torch.float32) |
| masks = torch.zeros(batch, max_tokens, dtype=torch.bool) |
| for row, sample in enumerate(samples): |
| length = sample["text_embeddings"].shape[0] |
| embeddings[row, :length] = sample["text_embeddings"] |
| masks[row, :length] = sample["text_mask"] |
| result = { |
| "latents": torch.stack([sample["latents"] for sample in samples]), |
| "text_embeddings": embeddings, |
| "text_mask": masks, |
| } |
| teacher_presence = ["teacher_features" in sample for sample in samples] |
| if any(teacher_presence) and not all(teacher_presence): |
| raise ValueError("all samples in a batch must either have teacher features or omit them") |
| if all(teacher_presence): |
| teacher_shapes = {tuple(sample["teacher_features"].shape) for sample in samples} |
| if len(teacher_shapes) != 1: |
| raise ValueError("teacher feature shapes must match within a batch") |
| result["teacher_features"] = torch.stack([sample["teacher_features"] for sample in samples]) |
| return result |
|
|
|
|
| def apply_classifier_free_dropout( |
| text_embeddings: Tensor, |
| text_mask: Tensor, |
| empty_embeddings: Tensor, |
| empty_mask: Tensor | None, |
| probability: float, |
| ) -> tuple[Tensor, Tensor]: |
| """Replace a random fraction of conditions with pre-encoded empty prompts.""" |
| if not 0 <= probability <= 1: |
| raise ValueError("probability must be between 0 and 1") |
| if probability == 0: |
| return text_embeddings, text_mask |
| batch, current_length, dim = text_embeddings.shape |
| if empty_embeddings.ndim == 2: |
| empty_embeddings = empty_embeddings.unsqueeze(0) |
| if empty_embeddings.ndim != 3 or empty_embeddings.shape[-1] != dim: |
| raise ValueError("empty_embeddings must have shape [L,D] or [1,L,D]") |
| if empty_embeddings.shape[0] not in (1, batch): |
| raise ValueError("empty_embeddings batch dimension must be 1 or match the training batch") |
| if empty_mask is None: |
| empty_mask = torch.ones(empty_embeddings.shape[:2], device=empty_embeddings.device, dtype=torch.bool) |
| elif empty_mask.ndim == 1: |
| empty_mask = empty_mask.unsqueeze(0) |
| if empty_mask.shape != empty_embeddings.shape[:2]: |
| raise ValueError("empty_mask must match the first two empty_embeddings dimensions") |
| max_length = max(current_length, empty_embeddings.shape[1]) |
| if empty_embeddings.shape[0] == 1: |
| empty_embeddings = empty_embeddings.expand(batch, -1, -1) |
| if empty_mask.shape[0] == 1: |
| empty_mask = empty_mask.expand(batch, -1) |
| expanded_text = text_embeddings.new_zeros(batch, max_length, dim) |
| expanded_mask = text_mask.new_zeros(batch, max_length) |
| expanded_text[:, :current_length] = text_embeddings |
| expanded_mask[:, :current_length] = text_mask |
| empty_embeddings = empty_embeddings.to(device=text_embeddings.device, dtype=text_embeddings.dtype) |
| empty_mask = empty_mask.to(device=text_mask.device, dtype=text_mask.dtype) |
| selected = torch.rand(batch, device=text_embeddings.device) < probability |
| expanded_text[:, : empty_embeddings.shape[1]] = torch.where( |
| selected[:, None, None], empty_embeddings, expanded_text[:, : empty_embeddings.shape[1]] |
| ) |
| expanded_mask[:, : empty_mask.shape[1]] = torch.where( |
| selected[:, None], empty_mask, expanded_mask[:, : empty_mask.shape[1]] |
| ) |
| return expanded_text, expanded_mask |
|
|
