Spaces:
Running on Zero
Running on Zero
| """Weighted sample-level mixing for multiple latent WebDataset sources.""" | |
| from __future__ import annotations | |
| import logging | |
| import math | |
| from collections.abc import Iterator, Mapping, Sequence | |
| from dataclasses import dataclass | |
| from typing import Any | |
| import torch | |
| from torch.utils.data import IterableDataset | |
| from .tensor import PytorchTensorWebDataset | |
| from .utils import ( | |
| DEFAULT_WDS_SAMPLE_SHUFFLE, | |
| DEFAULT_WDS_SAMPLE_SHUFFLE_INITIAL, | |
| derive_webdataset_seed, | |
| ) | |
| logger = logging.getLogger(__name__) | |
| _WDS_SOURCE_CHOICE_STREAM_ID = 0x5744_5305 | |
| _WDS_SOURCE_NAMESPACE_STRIDE = 0x100 | |
| class WebDatasetSource: | |
| """Validated configuration for one independently indexed WebDataset.""" | |
| source_id: str | |
| data_path: str | |
| annotation_path: str | |
| weight: float | |
| def normalize_webdataset_sources(raw_sources: Any) -> list[WebDatasetSource]: | |
| """Validate the public ``data_sources`` config and normalize its weights.""" | |
| if isinstance(raw_sources, (str, bytes)) or not isinstance(raw_sources, Sequence): | |
| raise TypeError("data_sources must be a non-empty list of mappings.") | |
| if not raw_sources: | |
| raise ValueError("data_sources must contain at least one WebDataset source.") | |
| sources: list[WebDatasetSource] = [] | |
| seen_ids: set[str] = set() | |
| for index, raw_source in enumerate(raw_sources): | |
| if not isinstance(raw_source, Mapping): | |
| raise TypeError(f"data_sources[{index}] must be a mapping.") | |
| source_id = raw_source.get("id", f"source_{index}") | |
| if not isinstance(source_id, str) or not source_id.strip(): | |
| raise ValueError(f"data_sources[{index}].id must be a non-empty string.") | |
| source_id = source_id.strip() | |
| if source_id in seen_ids: | |
| raise ValueError(f"Duplicate data source id: {source_id!r}.") | |
| seen_ids.add(source_id) | |
| data_path = raw_source.get("data_path") | |
| if not isinstance(data_path, str) or not data_path: | |
| raise ValueError(f"data_sources[{index}].data_path must be a non-empty string.") | |
| if not data_path.rstrip("/\\").lower().endswith("shards"): | |
| raise ValueError( | |
| f"data_sources[{index}].data_path must name a directory ending " | |
| f"in 'shards', got {data_path!r}." | |
| ) | |
| annotation_path = raw_source.get("annotation_path") | |
| if not isinstance(annotation_path, str) or not annotation_path: | |
| raise ValueError(f"data_sources[{index}].annotation_path must be a non-empty string.") | |
| weight = raw_source.get("weight") | |
| if isinstance(weight, bool) or not isinstance(weight, (int, float)): | |
| raise TypeError(f"data_sources[{index}].weight must be a positive number.") | |
| weight = float(weight) | |
| if not math.isfinite(weight) or weight <= 0: | |
| raise ValueError(f"data_sources[{index}].weight must be finite and greater than zero.") | |
| sources.append( | |
| WebDatasetSource( | |
| source_id=source_id, | |
| data_path=data_path, | |
| annotation_path=annotation_path, | |
| weight=weight, | |
| ) | |
| ) | |
| weight_sum = sum(source.weight for source in sources) | |
| return [ | |
| WebDatasetSource( | |
| source_id=source.source_id, | |
| data_path=source.data_path, | |
| annotation_path=source.annotation_path, | |
| weight=source.weight / weight_sum, | |
| ) | |
| for source in sources | |
| ] | |
| def validate_multi_source_selection( | |
| raw_sources: Any, | |
| *, | |
| legacy_fields: Mapping[str, Any], | |
| ) -> list[WebDatasetSource]: | |
| """Validate opt-in selection and reject ambiguous legacy fields.""" | |
| conflicting_fields = [name for name, value in legacy_fields.items() if value is not None] | |
| if conflicting_fields: | |
| joined = ", ".join(conflicting_fields) | |
| raise ValueError( | |
| f"data_sources cannot be combined with non-null legacy field(s): {joined}." | |
| ) | |
| return normalize_webdataset_sources(raw_sources) | |
| class WeightedMultiSourceDataset(IterableDataset): | |
| """Mix independent iterable datasets at sample granularity. | |
| Training sources are expected to be infinite resampled streams. For finite | |
| validation sources, an exhausted source is removed from the choice set and | |
| iteration ends after every source has been consumed once. | |
| """ | |
| def __init__( | |
| self, | |
| datasets: Sequence[IterableDataset], | |
| sources: Sequence[WebDatasetSource], | |
| *, | |
| seed: int, | |
| global_rank: int, | |
| resampled: bool, | |
| length: int | None = None, | |
| ) -> None: | |
| super().__init__() | |
| if len(datasets) != len(sources) or not datasets: | |
| raise ValueError("datasets and sources must have the same non-zero length.") | |
| if isinstance(seed, bool) or not isinstance(seed, int): | |
| raise TypeError("Multi-source seed must be a non-boolean integer.") | |
| if isinstance(global_rank, bool) or not isinstance(global_rank, int): | |
| raise TypeError("Multi-source global_rank must be a non-boolean integer.") | |
| if global_rank < 0: | |
| raise ValueError("Multi-source global_rank must be non-negative.") | |
| self.datasets = list(datasets) | |
| self.sources = list(sources) | |
| self._weights = [source.weight for source in sources] | |
| self._deterministic = True | |
| self._seed = seed | |
| self._global_rank = global_rank | |
| self._resampled = resampled | |
| if length is None: | |
| self._length = sum(len(dataset) for dataset in self.datasets) | |
| else: | |
| self._length = min( | |
| int(length), | |
| sum(len(dataset) for dataset in self.datasets), | |
| ) | |
| logger.info( | |
| "Initialized %s with sources=%s, resampled=%s, deterministic=%s, reported_length=%d", | |
| type(self).__name__, | |
| [ | |
| { | |
| "id": source.source_id, | |
| "weight": source.weight, | |
| "samples": len(dataset), | |
| } | |
| for source, dataset in zip( | |
| self.sources, | |
| self.datasets, | |
| strict=True, | |
| ) | |
| ], | |
| self._resampled, | |
| self._deterministic, | |
| self._length, | |
| ) | |
| def _worker_address(self) -> tuple[int, int]: | |
| worker = torch.utils.data.get_worker_info() | |
| if worker is None: | |
| return 0, 1 | |
| return worker.id, worker.num_workers | |
| def _choice_seed(self) -> int: | |
| worker_id, num_workers = self._worker_address() | |
| return derive_webdataset_seed( | |
| self._seed, | |
| self._global_rank, | |
| worker_id, | |
| num_workers, | |
| _WDS_SOURCE_CHOICE_STREAM_ID, | |
| ) | |
| def dataloader_seed(self) -> int | None: | |
| """Return an isolated DataLoader seed for deterministic mixing.""" | |
| return derive_webdataset_seed( | |
| self._seed, | |
| self._global_rank, | |
| worker_id=0, | |
| num_workers=1, | |
| stream_id=_WDS_SOURCE_CHOICE_STREAM_ID + 1, | |
| ) | |
| def _make_generator(self) -> torch.Generator: | |
| generator = torch.Generator(device="cpu") | |
| generator.manual_seed(self._choice_seed()) | |
| return generator | |
| def _choose_source( | |
| self, | |
| generator: torch.Generator, | |
| active: Sequence[bool], | |
| ) -> int: | |
| total_weight = sum( | |
| weight for weight, is_active in zip(self._weights, active, strict=True) if is_active | |
| ) | |
| if total_weight <= 0: | |
| raise RuntimeError("No active WebDataset source remains.") | |
| target = torch.rand((), generator=generator, dtype=torch.float64).item() * total_weight | |
| cumulative = 0.0 | |
| last_active = -1 | |
| for index, (weight, is_active) in enumerate(zip(self._weights, active, strict=True)): | |
| if not is_active: | |
| continue | |
| last_active = index | |
| cumulative += weight | |
| if target < cumulative: | |
| return index | |
| return last_active | |
| def _tag_sample(self, sample: dict[str, Any], source_index: int) -> dict[str, Any]: | |
| if not isinstance(sample, dict): | |
| raise TypeError( | |
| f"WebDataset source {self.sources[source_index].source_id!r} " | |
| f"yielded {type(sample).__name__}, expected dict." | |
| ) | |
| tagged = dict(sample) | |
| metadata = dict(tagged.get("metadata", {})) | |
| metadata.update( | |
| { | |
| "source_id": self.sources[source_index].source_id, | |
| "source_index": source_index, | |
| "source_weight": self.sources[source_index].weight, | |
| } | |
| ) | |
| tagged["metadata"] = metadata | |
| return tagged | |
| def __iter__(self) -> Iterator[dict[str, Any]]: | |
| iterators: list[Iterator[dict[str, Any]]] = [iter(dataset) for dataset in self.datasets] | |
| active = [True] * len(iterators) | |
| generator = self._make_generator() | |
| while any(active): | |
| source_index = self._choose_source(generator, active) | |
| try: | |
| sample = next(iterators[source_index]) | |
| except StopIteration: | |
| if self._resampled: | |
| source_id = self.sources[source_index].source_id | |
| raise RuntimeError( | |
| f"Resampled WebDataset source {source_id!r} exhausted unexpectedly." | |
| ) from None | |
| active[source_index] = False | |
| continue | |
| yield self._tag_sample(sample, source_index) | |
| def __len__(self) -> int: | |
| return self._length | |
| def build_multi_source_tensor_webdataset( | |
| raw_sources: Any, | |
| *, | |
| is_train: bool, | |
| seed: int, | |
| global_rank: int, | |
| length: int | None, | |
| target_frame_num: int | None = None, | |
| sample_shuffle: int = DEFAULT_WDS_SAMPLE_SHUFFLE, | |
| sample_shuffle_initial: int = DEFAULT_WDS_SAMPLE_SHUFFLE_INITIAL, | |
| ) -> WeightedMultiSourceDataset: | |
| """Build one latent WebDataset pipeline per source plus a weighted mixer.""" | |
| sources = normalize_webdataset_sources(raw_sources) | |
| datasets: list[PytorchTensorWebDataset] = [] | |
| for source_index, source in enumerate(sources): | |
| datasets.append( | |
| PytorchTensorWebDataset( | |
| data_path=source.data_path, | |
| annotation_path=source.annotation_path, | |
| target_frame_num=target_frame_num, | |
| shuffle=sample_shuffle if is_train else 0, | |
| shuffle_initial=sample_shuffle_initial, | |
| shardshuffle=is_train, | |
| resampled=is_train, | |
| seed=seed, | |
| global_rank=global_rank, | |
| seed_namespace=(source_index + 1) * _WDS_SOURCE_NAMESPACE_STRIDE, | |
| ) | |
| ) | |
| return WeightedMultiSourceDataset( | |
| datasets, | |
| sources, | |
| seed=seed, | |
| global_rank=global_rank, | |
| resampled=is_train, | |
| length=length, | |
| ) | |