# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 """Preprocess dataset for inference.""" from pathlib import Path from typing import Annotated, Any, Literal, cast, get_args import hydra import pydantic import tyro from loguru import logger from omegaconf import OmegaConf from torch.utils.data import Dataset, IterableDataset from typing_extensions import assert_never from cosmos_framework.inference.common.args import ( ConfigFileType, ResolvedFilePath, ResolvedPath, SetupArgs, ) from cosmos_framework.inference.common.config import deserialize_config_dict from cosmos_framework.inference.dataset_samples import ( DatasetSamples, IterableDatasetSamples, MapDatasetSamples, ) from cosmos_framework.scripts.dataset_utils import ( DEFAULT_GCS_CREDENTIALS, download_from_gcs_uri, override_dataset_root, ) Split = Literal["train", "val", "full"] ModelMode = Literal["forward_dynamics", "inverse_dynamics", "policy"] ALL_MODEL_MODES: list[str] = list(get_args(ModelMode)) def _select_dataloader( dataloaders_dict: dict[str, Any], dataset_name: str, ) -> tuple[str, dict[str, Any], dict[str, Any], Any]: """Pick the right dataloader entry from a multi-dataloader config. Selection strategy: 1. If *dataset_name* exactly matches a key in *dataloaders_dict*, use it. 2. Otherwise, search each dataloader's ``list_of_datasets`` for an entry whose ``name`` matches *dataset_name* and use the first hit. 3. Fall back to the first dataloader in iteration order. Returns ``(dataset_key, dataloader_config, dataset_config, inner_ds)`` where *dataset_config* is a mutable plain dict (via ``OmegaConf.to_container``) and *inner_ds* is the (possibly narrowed) ``list_of_datasets`` value. """ def _resolve(key: str) -> tuple[str, dict[str, Any], dict[str, Any], Any]: dl_cfg: dict[str, Any] = dataloaders_dict[key]["dataloader"] raw_ds = dl_cfg["dataset"] if not OmegaConf.is_config(raw_ds): raw_ds = OmegaConf.create(raw_ds) ds_cfg = cast(dict[str, Any], OmegaConf.to_container(raw_ds, resolve=True)) return key, dl_cfg, ds_cfg, ds_cfg.get("list_of_datasets") if dataset_name and dataset_name in dataloaders_dict: key, dl_cfg, ds_cfg, inner = _resolve(dataset_name) return key, dl_cfg, ds_cfg, inner if dataset_name: for candidate_key in dataloaders_dict: _, dl_cfg, ds_cfg, inner = _resolve(candidate_key) if inner is not None and isinstance(inner, list): names = [e.get("name") for e in inner if isinstance(e, dict)] if dataset_name in names: matched = [e for e in inner if isinstance(e, dict) and e.get("name") == dataset_name] ds_cfg["list_of_datasets"] = [matched[0]] inner = ds_cfg["list_of_datasets"] logger.info(f"Narrowed list_of_datasets to single entry: {dataset_name!r}") return candidate_key, dl_cfg, ds_cfg, inner available: dict[str, list[str | None]] = {} for k in dataloaders_dict: _, _, ds, lod = _resolve(k) if lod is not None and isinstance(lod, list): available[k] = [e.get("name") for e in lod if isinstance(e, dict)] else: available[k] = [] raise ValueError( f"dataset={dataset_name!r} not found in any dataloader. " f"Available dataloaders and their datasets: {available}" ) key, dl_cfg, ds_cfg, inner = _resolve(next(iter(dataloaders_dict))) return key, dl_cfg, ds_cfg, inner class DatasetArgs(pydantic.BaseModel): """Arguments controlling which dataset to load and how to prepare it. The Hydra ``config_file`` and ``experiment`` used to resolve the dataset are supplied separately (typically from the setup args) so that CLI flag names never collide with the model's setup arguments. """ model_config = pydantic.ConfigDict(extra="forbid", use_attribute_docstrings=True) dataset_experiment: str = "" """Hydra experiment for the dataset config. If empty, the setup experiment is used.""" dataset: str = "" """Dataset selector (matched against dataloader keys or list_of_datasets names).""" sample_name: str = "" """Override the dataset name used in sample paths (output directories, ground truth). Falls back to ``dataset`` if empty.""" dataset_split: Split = "val" """Dataset split.""" model_mode: ModelMode | Literal["joint", "vision"] = "joint" """Model mode(s) to iterate / dispatch on. ``joint`` expands to all three action modes (forward_dynamics / inverse_dynamics / policy). ``vision`` dispatches the eval script to score pre-generated videos against a GT directory (no inference).""" num_samples: int | None = None """Maximum number of samples to load.""" sample_stride: int = 1 """Stride between sample indices.""" gcs_path_map: dict[str, str] = {} """Maps default dataset root paths to S3 URIs. When the resolved config root(s) matches a key, the data is downloaded from the corresponding S3 URI and the root is overridden with the local cache path.""" gcs_credentials: str = DEFAULT_GCS_CREDENTIALS """Path to GCS credentials JSON file.""" root_override: str = "" """Override the dataset root path without downloading. Takes precedence over gcs_path_map.""" gcs_root_override: str = "" """S3 URI to download and use as the dataset root. Takes precedence over gcs_path_map but unlike root_override, triggers a download first.""" force_download: bool = False """Force re-download even if local dataset directory already exists.""" dataset_name_override: str = "" """Override the inner dataset's name field after narrowing.""" dataset_kwargs: dict[str, Any] = {} """Extra keyword arguments applied to the inner dataset config before instantiation. Useful for overriding constructor parameters like ``snap_to_subtask`` or ``max_subtasks_per_episode``.""" cache_dir: Annotated[ResolvedPath | None, tyro.conf.arg(aliases=("-c",))] = None """Local cache directory for GCS downloads. Required when ``gcs_path_map`` triggers a download.""" config_file: ResolvedFilePath | None = None """Path to a serialized YAML dataset config file.""" @property def dataset_label(self) -> str: """Label used for sample output paths; falls back to the selector when unset.""" return self.sample_name or self.dataset def create_dataset( args: DatasetArgs, config_args: SetupArgs | None = None, ) -> DatasetSamples: """Load a dataset and return a lazy ``DatasetSamples`` iterator of ``(sample_args, data_batch)`` pairs. Two config sources are supported (mirroring checkpoint handling): * **config store:** provide ``config_args`` with ``config_file`` + ``experiment`` to resolve the config via Hydra. * **YAML:** set ``args.dataset_yaml`` to a serialized YAML config Args: args: Dataset selection and download parameters. config_args: Hydra config arguments (internal config store). """ # --- load config ---------------------------------------------------------- if args.config_file is not None: config = deserialize_config_dict(args.config_file) elif config_args is not None: if args.dataset_experiment: config_args = config_args.model_copy(update={"experiment": args.dataset_experiment}) # load_config() only handles .py module configs; YAML/JSON use the dict path. if config_args.config_file_type == ConfigFileType.MODULE: config = config_args.load_config() else: config = deserialize_config_dict(Path(config_args.config_file)) else: raise ValueError("Provide 'config_args' or set 'dataset_yaml'") def _cfg(key: str) -> Any: return config[key] if isinstance(config, dict) else getattr(config, key) _train = _cfg("dataloader_train") _val = _cfg("dataloader_val") if args.dataset_split == "train": dataloaders_config = _train elif args.dataset_split in ("val", "full"): dataloaders_config = _val if _val is not None else _train else: assert_never(args.dataset_split) assert dataloaders_config is not None, "Neither dataloader_val nor dataloader_train found in config" dataloaders_dict: dict[str, Any] = dataloaders_config["dataloaders"] dataset_key, _, dataset_config, inner_ds = _select_dataloader(dataloaders_dict, args.dataset) logger.info(f"Selected dataloader {dataset_key!r} for dataset={args.dataset!r}") if isinstance(inner_ds, list): if len(inner_ds) == 1 and isinstance(inner_ds[0], dict): inner_ds = inner_ds[0] else: raise ValueError("Use --dataset to select a single dataset entry by name.") resolution: str | None = inner_ds.get("resolution") if isinstance(inner_ds, dict) else None # wrap_dataset also accepts a global resolution that individual entries can override if resolution is None: resolution: str | None = dataset_config.get("resolution") # --- resolve the inner dataset target ------------------------------------- # Dataset config fields may be nested inside inner_ds["dataset"] when the # full dataset_entry wrapper dict is preserved (after _select_dataloader narrowing), # or live directly on inner_ds when the raw dataset dict is returned. _inner = inner_ds if isinstance(inner_ds, dict) and isinstance(inner_ds.get("dataset"), dict): _inner = inner_ds["dataset"] # --- apply dataset_name override ------------------------------------------ if args.dataset_name_override and isinstance(_inner, dict): for key in ("dataset_name", "name"): if key in _inner: logger.info(f"Overriding inner dataset {key}={_inner[key]!r} -> {args.dataset_name_override!r}") _inner[key] = args.dataset_name_override break # --- resolve root via overrides / gcs_path_map and download ---------- root_override: str | list[str] = "" config_root = _inner.get("root") if isinstance(_inner, dict) else None if args.root_override: root_override = args.root_override elif args.gcs_root_override: if args.cache_dir is None: raise ValueError("cache_dir is required when gcs_root_override triggers a download") args.cache_dir.mkdir(parents=True, exist_ok=True) root_override = str( download_from_gcs_uri( args.gcs_root_override, args.cache_dir, gcs_credentials=args.gcs_credentials, force_download=args.force_download, ) ) logger.info(f"Downloaded dataset from {args.gcs_root_override} to {root_override}") elif (config_root is not None) and args.gcs_path_map: root_is_sequence = False if isinstance(config_root, str): orig_roots: list[str] = [config_root] elif isinstance(config_root, (list, tuple)): root_is_sequence = True orig_roots = list(config_root) else: raise ValueError(f"Unexpected root type: {type(config_root)}") # Populate the cache directory. if args.cache_dir is None: raise ValueError("cache_dir is required when gcs_path_map triggers a download") args.cache_dir.mkdir(parents=True, exist_ok=True) for i, orig_root in enumerate(orig_roots): if orig_root.startswith("s3://"): # Already an S3 URI — use directly, no mapping needed. logger.info(f"Root {orig_root} is already an S3 URI, skipping gcs_path_map") continue uri = args.gcs_path_map.get(orig_root.rstrip("/")) if uri is None: raise ValueError(f"No entry found in gcs_path_map for root='{orig_root.rstrip('/')}'") orig_roots[i] = str( download_from_gcs_uri( uri, args.cache_dir, gcs_credentials=args.gcs_credentials, force_download=args.force_download, dataset_name_override=args.dataset_name_override, ) ) if orig_roots: root_override = orig_roots if root_is_sequence else orig_roots[0] if root_override: override_dataset_root(dataset_config, root_override) # --- override split ------------------------------------------------------- if isinstance(_inner, dict) and "split" in _inner: _inner["split"] = args.dataset_split logger.info(f"Overriding dataset split to {args.dataset_split!r}") if isinstance(_inner, dict) and "is_val" in _inner: is_val = args.dataset_split == "val" _inner["is_val"] = is_val logger.info(f"Overriding dataset is_val to {is_val!r}") # --- disable training-only augmentations for val and 'full' -------------------------- if args.dataset_split in ("val", "full") and "cfg_dropout_rate" in dataset_config: dataset_config["cfg_dropout_rate"] = 0.0 logger.info("Overriding cfg_dropout_rate to 0.0 for val export") # --- disable shuffle for val dataset export ------------------------------- if args.dataset_split == "val" and isinstance(_inner, dict) and "shuffle" in _inner: _inner["shuffle"] = False logger.info("Overriding dataset shuffle to False for val export") # --- apply dataset_kwargs overrides -------------------------------------- if args.dataset_kwargs and isinstance(_inner, dict): for key, value in args.dataset_kwargs.items(): logger.info(f"Overriding dataset {key}={_inner.get(key)!r} -> {value!r}") _inner[key] = value # --- apply gcs credentials override (internal-only)-------------------- if isinstance(_inner, dict) and "credential_path" in _inner: _inner["credential_path"] = args.gcs_credentials logger.info(f"Overriding dataset credential_path to {_inner['credential_path']!r}") # --- instantiate dataset -------------------------------------------------- logger.info(f"Instantiating dataset with config: {dataset_config}") dataset: Dataset = hydra.utils.instantiate(dataset_config) dataset_size = len(dataset) # type: ignore transform = getattr(dataset, "transform", None) if isinstance(dataset, IterableDataset): inner = getattr(dataset, "dataset", None) raw = getattr(inner, "dataset", inner) if raw is not None and hasattr(raw, "__getitem__"): dataset = raw if hasattr(dataset, "_register_sources") and hasattr(dataset, "_all_shard_roots") and dataset_size == 0: dataset._register_sources() dataset_size = len(dataset) # type: ignore def _create_dataset_samples(ds: Dataset, m: list[str], ids: list[int]) -> DatasetSamples: sample_overrides_data = config_args.sample_overrides.model_dump(exclude_none=True) if config_args else {} if isinstance(ds, IterableDataset): return IterableDatasetSamples(ds, m, ids, transform, resolution, args.dataset_label, sample_overrides_data) return MapDatasetSamples(ds, m, ids, transform, resolution, args.dataset_label, sample_overrides_data) # --- compute sample indices ----------------------------------------------- sample_ids = list(range(0, dataset_size, args.sample_stride)) if args.num_samples: sample_ids = sample_ids[: args.num_samples] modes = list(ALL_MODEL_MODES) if args.model_mode == "joint" else [args.model_mode] # The transform's text tokenizer produces text_token_ids for training only; # during inference the model re-tokenizes from the caption string (see # OmniMoTModel._get_inference_text_tokens). Disable it to avoid tokenizer # compatibility issues in environments where apply_chat_template returns a dict. if transform is not None and hasattr(transform, "text_tokenizer"): transform.text_tokenizer = None logger.info(f"Created lazy iterator for {len(modes) * len(sample_ids)} samples across modes {modes}") return _create_dataset_samples(dataset, modes, sample_ids)