Spaces:
Running on L40S
Running on L40S
File size: 6,434 Bytes
9f818c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
"""Lazy dataset sample iterators for map-style and iterable-style datasets."""
import itertools
import json
from collections.abc import Iterable, Iterator
from typing import Any, Callable
import torch
from loguru import logger
from torch.utils.data import Dataset, IterableDataset
from torch.utils.data.dataloader import default_collate
from cosmos_framework.inference.args import OmniSampleOverrides
from cosmos_framework.scripts.dataset_utils import set_dataset_mode
from cosmos_framework.utils.vfm.data_utils import get_vision_data_resolution
Sample = tuple[OmniSampleOverrides, dict[str, Any]]
def _collate_sample(sample: dict) -> dict:
"""Collate a single sample dict, adding a batch dim to tensors."""
result: dict[str, Any] = {}
for key, val in sample.items():
if isinstance(val, torch.Tensor):
result[key] = val.unsqueeze(0)
else:
try:
result[key] = default_collate([val])
except TypeError:
result[key] = [val]
return result
def _normalize_caption(raw_sample: dict) -> str:
"""Normalize ``ai_caption`` to a string in-place and return it.
JSON-dict captions (from ``ActionPromptJsonFormatter``) are serialized
so collation, batch merging, and the model input treat them identically
to plain-text captions, matching the training side's
``TextTokenizerTransform``.
Raises:
TypeError: If ``ai_caption`` is present and is neither ``str`` nor
``dict``.
"""
caption = raw_sample.get("ai_caption", "")
if isinstance(caption, dict):
caption = json.dumps(caption)
raw_sample["ai_caption"] = caption
elif not isinstance(caption, str):
raise TypeError(f"ai_caption must be str or dict, got {type(caption).__name__}")
return caption
class _BaseSamples(Iterable[Sample]):
"""Base iterator yielding ``(OmniSampleOverrides, data_batch)`` pairs.
Iterates over every (mode, sample_id) combination, applying an optional
transform to each raw item. Subclasses implement ``__iter__`` for
map-style and iterable-style datasets respectively.
"""
def __init__(
self,
dataset: Dataset | IterableDataset,
modes: list[str], # model modes to iterate (e.g. ["joint", "forward_dynamics", etc.])
sample_ids: list[int], # indices into dataset to yield
transform: Callable | None, # UVA transform pipeline applied per item, or None
resolution: str | None, # global resolution override; inferred from video shape if None
dataset_name: str, # name of the dataset"
sample_overrides_data: dict[str, Any] | None = None, # additional overrides to apply to every sample
) -> None:
self._dataset = dataset
self._modes = modes
self._sample_ids = sample_ids
self._transform = transform
self._resolution = resolution
self._dataset_name = dataset_name
self._sample_overrides_data = sample_overrides_data
def __len__(self) -> int:
return len(self._modes) * len(self._sample_ids)
def _make_sample_from_raw(self, raw_sample: Any, sample_idx: int, mode: str) -> Sample:
"""Apply transform, collate, and wrap a raw dataset item into a ``Sample``."""
resolution = self._resolution
if resolution is None:
video = raw_sample.get("video")
if video is not None:
resolution = get_vision_data_resolution(video.shape[-2:])
if self._transform is not None:
raw_sample = self._transform(raw_sample, resolution=resolution)
prompt = _normalize_caption(raw_sample)
sample_data = _collate_sample(raw_sample)
sample_name = f"{self._dataset_name}/{mode}/{sample_idx}" if self._dataset_name else f"{mode}/{sample_idx}"
sample_args = OmniSampleOverrides(
name=sample_name,
prompt=prompt,
resolution=resolution, # type: ignore
raw_action_dim=sample_data.get("raw_action_dim", [None])[0],
)
# Apply any additional sample overrides specified in the setup config (e.g. num_steps, guidance, etc.)
sample_args = sample_args.model_copy(update=self._sample_overrides_data)
return sample_args, sample_data
class MapDatasetSamples(_BaseSamples):
"""Iterator for map-style datasets (``Dataset``), accessed via ``__getitem__``.
Iterates modes in order, indexing each sample directly by its — enabling
random access.
"""
def __iter__(self) -> Iterator[Sample]:
for mode in self._modes:
set_dataset_mode(self._dataset, mode)
for sample_idx in self._sample_ids:
raw_sample = self._dataset[sample_idx] # type: ignore[index]
yield self._make_sample_from_raw(raw_sample, sample_idx, mode)
class IterableDatasetSamples(_BaseSamples):
"""Iterator for iterable-style datasets (``IterableDataset``), accessed via ``__iter__``.
Since random access is unavailable, advances the underlying iterator using
``islice`` to reach each target sample index in order — requires
``sample_ids`` to be sorted ascending.
"""
def __iter__(self) -> Iterator[Sample]:
for mode in self._modes:
set_dataset_mode(self._dataset, mode)
dataset = iter(getattr(self._dataset, "dataset", self._dataset))
cur_ix = 0
for sample_idx in sorted(self._sample_ids):
try:
raw_sample = next(itertools.islice(dataset, sample_idx - cur_ix, None))
except StopIteration:
# Dataset exhausted early (inaccurate __len__); move on to next mode.
logger.warning(
f"Dataset {self._dataset_name!r} exhausted early while iterating mode={mode!r}: "
f"tried to reach sample_idx={sample_idx}, expected __len__={len(self._dataset)}. " # type: ignore[arg-type]
"Moving on to next mode."
)
break
cur_ix = sample_idx + 1
yield self._make_sample_from_raw(raw_sample, sample_idx, mode)
DatasetSamples = MapDatasetSamples | IterableDatasetSamples
|