# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 import hashlib import json import pickle from collections.abc import Callable, Generator, Iterable from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, Sequence, TypeVar, cast, override import attrs import cattrs import cattrs.preconf.json import safetensors.torch import torch from PIL import Image from torch.utils._pytree import tree_map_only from torch.utils.data import Dataset from typing_extensions import Self from cosmos_framework.inference.args import ( ModelMode, NegativeMetadataMode, OmniSampleArgs, OmniSetupArgs, ) from cosmos_framework.inference.common.args import ( CheckpointType, ConfigFileType, ParallelismArgs, SampleArgs, SampleOutput, SampleOutputs, SetupArgs, ) from cosmos_framework.inference.common.inference import Inference, sync_distributed_errors from cosmos_framework.inference.common.init import get_rank, get_world_size from cosmos_framework.inference.model import Cosmos3OmniConfig, Cosmos3OmniModel from cosmos_framework.inference.vision import ( build_conditioned_video_batch, build_image_edit_batch, load_conditioning_image, load_conditioning_video, load_prompt_upsampling_image, pil_to_conditioning_frames, resize_pil_image, ) from cosmos_framework.utils import log from cosmos_framework.tools.visualize.video import save_img_or_video from cosmos_framework.configs.base.defaults.compile import CompileConfig from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig from cosmos_framework.model.vfm.omni_mot_model import OmniMoTModel from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import _SYSTEM_PROMPT_IMAGE_EDITING from cosmos_framework.model.vfm.upsampler.prompts import is_upsampled_prompt if TYPE_CHECKING: from cosmos_framework.configs.base.defaults.model_config import OmniMoTModelConfig UpsampleTask = Literal["t2i", "t2v", "i2v"] _BatchItem = TypeVar("_BatchItem") def _iter_packed_batches( items: Iterable[_BatchItem], get_sample_args: Callable[[_BatchItem], OmniSampleArgs], model: OmniMoTModel, max_model_len: int | None, max_num_seqs: int | None, ) -> Generator[list[_BatchItem]]: """Greedily pack a stream of items into batches under a single budget. Walks ``items`` once, in order, and accumulates them into the current batch. Before adding an item, the helper checks whether doing so would exceed the configured budget; if it would (and the current batch is non-empty), the current batch is emitted and a fresh one is started with the item. If a single item exceeds the budget on its own, it still gets emitted in a batch by itself (no item is ever dropped). Order is preserved. The helper is item-agnostic: it never inspects ``items`` beyond passing each one to ``get_sample_args`` to get the ``OmniSampleArgs`` used for token counting. Callers can therefore batch indices, ``(args, data)`` pairs, or any other value type, and the yielded batches keep the same element type. Exactly one of ``max_model_len`` and ``max_num_seqs`` must be provided — that selects which budget is enforced: - ``max_model_len``: token budget. Each item's token count is computed via ``_compute_num_tokens_for_sample(sa, model)``. - ``max_num_seqs``: sequence-count budget. Each item counts as one slot. Streaming: at any time only the current in-flight batch is held in memory; ``items`` is consumed lazily. Args: items: Iterable of arbitrary items (typed as ``_BatchItem``) to pack. Consumed exactly once, in order. get_sample_args: Callable mapping each item to its ``OmniSampleArgs``. Used for token counting and for the per-item invariants below. model: The model, forwarded to ``_compute_num_tokens_for_sample`` for token-budget mode. Unused (but still required) for sequence-count mode. max_model_len: Token budget per batch. Mutually exclusive with ``max_num_seqs``; exactly one must be set (the other ``None``). max_num_seqs: Sequence-count budget per batch. Mutually exclusive with ``max_model_len``; exactly one must be set (the other ``None``). Yields: ``list[_BatchItem]``: a non-empty batch of items, in input order. The yielded list is owned by the caller — the helper drops its reference after yielding and starts fresh, so callers are free to mutate or store it. Raises: AssertionError: If neither or both of ``max_model_len`` / ``max_num_seqs`` are provided, if ``get_sample_args`` returns a non-``OmniSampleArgs``, or if any item's ``num_outputs != 1``. Callers must seed-expand multi-output samples (e.g. via ``_finalize_sample_args_list``) before passing them in. """ assert max_model_len is not None or max_num_seqs is not None, "Either max_model_len or max_num_seqs must be set" assert max_model_len is None or max_num_seqs is None, "Either max_model_len or max_num_seqs must be set, not both" cur: list[_BatchItem] = [] running_tokens = 0 running_seqs = 0 for item in items: sa = get_sample_args(item) assert isinstance(sa, OmniSampleArgs) assert sa.num_outputs == 1, "num_outputs must be 1" if max_model_len is not None: num_tokens = _compute_num_tokens_for_sample(sa, model) if cur and running_tokens + num_tokens > max_model_len: yield cur cur = [] running_tokens = 0 running_tokens += num_tokens else: if cur and running_seqs + 1 > max_num_seqs: # type: ignore[operator] yield cur cur = [] running_seqs = 0 running_seqs += 1 cur.append(item) if cur: yield cur def _fallback_seed(sample_args: OmniSampleArgs) -> int: """Stable per-sample fallback seed used when ``sample_args.seed is None``. We derive the seed deterministically from the sample's identity (``name`` + ``output_dir``) instead of calling ``random.randint``, because each rank of a CP / CFG-parallel replica must agree on the seed: the sampler uses it to draw the initial noise, and divergent noise across ranks corrupts the collective denoising loop. Independent calls to Python's global ``random`` module (which is what the previous code did) return different values on different ranks unless the user has separately seeded it identically everywhere, which is easy to forget. The returned int fits in 31 bits so it's safe to pass to any downstream API expecting a non-negative ``int32`` seed. """ identity = f"{sample_args.name}|{sample_args.output_dir}".encode("utf-8") return int.from_bytes(hashlib.sha256(identity).digest()[:4], "big") & 0x7FFFFFFF def _compute_num_tokens_for_sample(sample_args: OmniSampleArgs, model: OmniMoTModel) -> int: """Estimate the number of tokens for a single inference sample. Follows the counting logic in ``JointDataLoader._compute_num_tokens_per_sample`` (vision + text + EOS). """ w, h = sample_args.vision_size T = sample_args.num_frames spatial_cf = cast(int, model.tokenizer_vision_gen.spatial_compression_factor) temporal_cf = cast(int, model.tokenizer_vision_gen.temporal_compression_factor) patch_spatial: int = model.config.diffusion_expert_config.patch_spatial vae_spatial_downsample = spatial_cf * patch_spatial vae_temporal_downsample = temporal_cf latent_h = h // vae_spatial_downsample latent_w = w // vae_spatial_downsample latent_t = 1 + (T - 1) // vae_temporal_downsample num_vision_tokens = latent_h * latent_w * latent_t # small compared to vision tokens, so we can ignore them for now. return num_vision_tokens def _infer_native_prompt_upsampling_tasks( data_batch: dict[str, Any], sample_args_list: Sequence[OmniSampleArgs], model: OmniMoTModel, ) -> list[UpsampleTask | None]: """Return the per-sample native V4.2 prompt-upsample task list. Each entry corresponds positionally to ``sample_args_list[i]`` and is either the canonical V4.2 task (``"t2v"`` / ``"t2i"`` / ``"i2v"``) when native upsampling should run for that sample, or ``None`` when it should be skipped. Per-sample reasons for ``None`` include: * Action-only samples (``model_mode.is_action``) — the V4.2 templates are vision-caption only. * Samples that did not opt in via ``native_prompt_upsampling=True`` (see ``OmniSampleOverrides.build_sample`` for how this flag is derived from ``prompt_upsampling`` × ``prompt_upsampling_applied``). * Samples whose prompt already looks like upsampler output (fenced or bare JSON; see :func:`projects.cosmos3.vfm.upsampler.prompts.is_upsampled_prompt`). * Batch-wide modality ambiguity — ``data_batch`` tensors are stacked along the batch dim, so the image/video keys are inherently batch-uniform; if both keys are present (mixed batch) or neither is (no vision input), every sample is marked ``None``. * Image-editing samples (image input *plus* per-sample vision conditioning) — not yet supported by the V4.2 upsampler. * I2V samples when ``data_batch`` lacks the VLM-ready ``_prompt_upsampling_images`` side channel. * I2V samples when the reasoner LM was loaded with ``include_visual=False`` (the ViT is required to encode the conditioning frame for the V4.2 i2v template). Mixed batches — i.e. some samples returning a task and others ``None``, or different samples returning different tasks — are *allowed* by this function; the caller is responsible for routing them to the model. """ n_samples = len(sample_args_list) # Batch-wide modality presence: ``data_batch`` tensors are stacked # along the batch dim, so the image/video keys are inherently # batch-uniform. If both keys are present (mixed modality) or # neither is (no vision input), every sample is ambiguous. has_image = model.input_image_key in data_batch has_video = model.input_video_key in data_batch if has_image == has_video: return [None] * n_samples plans = data_batch.get("sequence_plan") or [None] * n_samples has_upsampling_images = "_prompt_upsampling_images" in data_batch # The reasoner's ViT is constructed only when the MoT wrapper was # built with ``include_visual=True`` (see ``Qwen3VLTextForCausalLM.__init__`` # and friends in ``unified_mot``). Without it, the i2v branch in # ``_impl_generate_reasoner_text`` raises ``ValueError`` at the # ``hasattr(causal_lm, "visual")`` gate — fall back to no-op # upsampling so the (already-captioned) prompt flows to diffusion # unchanged. Walk the ``model.net.language_model.visual`` chain # defensively with ``getattr(..., None)`` so test mocks that don't # set up the full reasoner-LM attribute chain (e.g. # ``tests/test_eval_model.py``'s mock for the # ``EvalModel → OmniInference.generate_batch`` boundary) gracefully # land on ``has_visual_tower = False`` instead of raising # ``AttributeError`` here. language_model = getattr(getattr(model, "net", None), "language_model", None) has_visual_tower = hasattr(language_model, "visual") results: list[UpsampleTask | None] = [] for i, sample_args in enumerate(sample_args_list): if sample_args.model_mode.is_action: results.append(None) continue if not sample_args.native_prompt_upsampling: results.append(None) continue # Content-based per-sample check: skip native prompt upsampling # when the prompt already looks like the V4.2 upsampler's output # (fenced ``json`` payload or bare JSON object). Re-running # upsampling on a JSON-shaped prompt would corrupt it. if is_upsampled_prompt(sample_args.prompt): results.append(None) continue plan = plans[i] if i < len(plans) else None sample_has_conditioning = bool(getattr(plan, "condition_frame_indexes_vision", [])) if has_image: # Image editing (image input + vision conditioning) is not # yet supported by the V4.2 upsampler templates. results.append(None if sample_has_conditioning else "t2i") continue # Video path. if not sample_has_conditioning: results.append("t2v") continue # I2V: requires both the VLM-ready side-channel images and a # reasoner with a visual tower. if not has_upsampling_images: raise ValueError( "I2V prompt upsampling requires '_prompt_upsampling_images' with one VLM-ready image per caption." ) if not has_visual_tower: raise ValueError( "I2V prompt upsampling requires the reasoner LM to have a visual tower (include_visual=True)" ) results.append("i2v") return results def _format_prompt_with_template( prompt: str, *, fps: int, num_frames: int, duration_template: str | None, resolution_template: str | None, h: int, w: int, force_duration_template: bool = False, ) -> str: """Append duration/fps and resolution metadata to a prompt.""" prompt = prompt.strip() if duration_template is not None and (num_frames > 1 or force_duration_template): duration = num_frames / fps dur_text = duration_template.format(duration=duration, fps=fps) prompt = prompt.rstrip(".") + ". " + dur_text prompt = prompt.strip() if resolution_template is not None: res_text = resolution_template.format(height=h, width=w) prompt = prompt.rstrip(".") + ". " + res_text return prompt def _parse_json_object_prompt(prompt: str) -> dict | None: """Return the parsed dict iff ``prompt`` is a JSON object string; else ``None``. JSON arrays / numbers / strings / nulls are NOT considered "JSON-object prompts" and return ``None`` so they continue down the plain-text path. """ try: obj = json.loads(prompt) except (json.JSONDecodeError, TypeError, ValueError): return None return obj if isinstance(obj, dict) else None def _format_json_prompt_with_template( prompt_obj: dict, *, fps: int, num_frames: int, aspect_ratio: str | None, h: int, w: int, include_temporal_metadata: bool, ) -> str: """JSON-prompt counterpart to ``_format_prompt_with_template``. Injects structured metadata fields directly into the parsed prompt object, matching the training-time augmentors so the tokenizer sees the exact schema the model was trained on: - ``ResolutionTextInfo`` -> ``resolution: {"H": int, "W": int}``, ``aspect_ratio: str`` - ``DurationFPSTextTimeStamps`` -> ``duration: "s"``, ``fps: float`` for video samples only Always overwrites existing values for these keys, mirroring the augmentors' ``dict.update(...)`` semantics: the actual generation specs are the source of truth, regardless of what the input prompt may have specified. """ metadata: dict[str, Any] = {} if include_temporal_metadata: duration_seconds = int(num_frames / fps) if fps > 0 else 0 metadata.update( { "duration": f"{duration_seconds}s", "fps": float(fps), } ) else: prompt_obj.pop("duration", None) prompt_obj.pop("fps", None) metadata["resolution"] = {"H": int(h), "W": int(w)} if aspect_ratio is not None: metadata["aspect_ratio"] = aspect_ratio prompt_obj.update(metadata) log.debug(f"Injected JSON-prompt metadata fields: {sorted(metadata.keys())}") return json.dumps(prompt_obj) def _get_prompt_sample_data(sample_args: OmniSampleArgs, model: OmniMoTModel, *, h: int, w: int, device: Any) -> dict: duration_template = sample_args.duration_template inverse_duration_template = sample_args.inverse_duration_template prompt_obj = _parse_json_object_prompt(sample_args.prompt) prompt_is_json = prompt_obj is not None if prompt_is_json: assert prompt_obj is not None # type-narrowing prompt = _format_json_prompt_with_template( prompt_obj, fps=sample_args.fps, num_frames=sample_args.num_frames, aspect_ratio=sample_args.aspect_ratio, h=h, w=w, include_temporal_metadata=sample_args.num_frames > 1, ) elif not sample_args.native_prompt_upsampling: prompt = _format_prompt_with_template( sample_args.prompt, fps=sample_args.fps, num_frames=sample_args.num_frames, duration_template=duration_template, resolution_template=sample_args.resolution_template, h=h, w=w, ) else: # If native prompt upsampling is enabled, the duration and resolution # metadata are added into the upsampled JSON prompt directly. prompt = sample_args.prompt.strip() out = { model.input_caption_key: [prompt] * sample_args.num_outputs, } negative_prompt = sample_args.negative_prompt if sample_args.negative_metadata_mode == NegativeMetadataMode.SAME: negative_prompt = ( _format_prompt_with_template( negative_prompt if negative_prompt is not None else "", fps=sample_args.fps, num_frames=sample_args.num_frames, duration_template=duration_template, resolution_template=sample_args.resolution_template, h=h, w=w, ) .lstrip(".") .strip() ) elif sample_args.negative_metadata_mode == NegativeMetadataMode.INVERSE: negative_prompt = ( _format_prompt_with_template( negative_prompt if negative_prompt is not None else "", fps=sample_args.fps, num_frames=sample_args.num_frames, duration_template=inverse_duration_template, resolution_template=sample_args.inverse_resolution_template, h=h, w=w, force_duration_template=True, ) .lstrip(".") .strip() ) if negative_prompt: neg_key = "neg_" + model.input_caption_key out[neg_key] = [negative_prompt] * sample_args.num_outputs return out def _get_reasoner_sample_data(sample_args: OmniSampleArgs, model: OmniMoTModel) -> dict[str, Any]: """Sample batch for reasoner text generation: prompt + optional conditioning image.""" image: Image.Image | None = None if sample_args.vision_path is not None: image = Image.open(sample_args.vision_path).convert("RGB") return { model.input_caption_key: [sample_args.prompt], "reasoner_images": [image], } def _get_image_edit_sample_data( sample_args: OmniSampleArgs, model: OmniMoTModel, *, device: Any, ) -> dict: """Create a sample batch for image-edit generation.""" assert sample_args.vision_path is not None if sample_args.resolution and sample_args.aspect_ratio: w, h = sample_args.vision_size conditioning_frames = load_conditioning_image(Path(sample_args.vision_path), target_h=h, target_w=w) else: pil_img = Image.open(sample_args.vision_path).convert("RGB") pil_img = resize_pil_image(pil_img, max_size=512, padding_constant=32) conditioning_frames, h, w = pil_to_conditioning_frames(pil_img) conditioning_frames = conditioning_frames.to(device=device) batch = build_image_edit_batch(conditioning_frames, h=h, w=w, batch_size=sample_args.num_outputs) batch["system_prompt"] = _SYSTEM_PROMPT_IMAGE_EDITING batch |= _get_prompt_sample_data(sample_args, model, h=h, w=w, device=device) return batch def get_sample_data( sample_args: OmniSampleArgs, model: OmniMoTModel, *, device: Any = "cuda", ) -> dict: """Create a sample batch for generation.""" if sample_args.model_mode.is_reasoner: return _get_reasoner_sample_data(sample_args, model) if sample_args.transfer_hints: return {} if sample_args.model_mode.is_action: from cosmos_framework.inference.action import get_action_sample_data assert sample_args.vision_path is not None return get_action_sample_data( model_config=model, batch_size=sample_args.num_outputs, prompt=sample_args.prompt, vision_path=sample_args.vision_path, model_mode=sample_args.model_mode, action_path=sample_args.action_path, domain_name=sample_args.domain_name, view_point=sample_args.view_point, resolution=str(sample_args.image_size), action_chunk_size=sample_args.action_chunk_size, max_action_dim=model.config.max_action_dim, fps=sample_args.fps, device=device, ) if sample_args.model_mode == ModelMode.IMAGE2IMAGE: return _get_image_edit_sample_data(sample_args, model, device=device) w, h = sample_args.vision_size if sample_args.num_frames == 1: input_vision_key = model.input_image_key else: input_vision_key = model.input_video_key with torch.device(device): prompt_upsampling_image: Image.Image | None = None match sample_args.condition_vision_mode: case "image": assert sample_args.vision_path is not None vision_path = Path(sample_args.vision_path) conditioning_frames = load_conditioning_image(vision_path, target_h=h, target_w=w) prompt_upsampling_image = load_prompt_upsampling_image(vision_path, target_h=h, target_w=w) case "video": assert sample_args.vision_path is not None assert sample_args.condition_frame_indexes_vision is not None num_condition_latent_frames = max(sample_args.condition_frame_indexes_vision) + 1 max_frames = model.tokenizer_vision_gen.get_pixel_num_frames(num_condition_latent_frames) conditioning_frames = load_conditioning_video( Path(sample_args.vision_path), target_h=h, target_w=w, max_frames=max_frames, keep=sample_args.condition_video_keep or "first", ) case _: conditioning_frames = None if conditioning_frames is not None: assert sample_args.condition_frame_indexes_vision is not None conditioned = build_conditioned_video_batch( conditioning_frames, condition_frames_vision=sample_args.condition_frame_indexes_vision, w=w, h=h, num_frames=sample_args.num_frames, fps=sample_args.fps, batch_size=sample_args.num_outputs, ) # Keep the list form (rather than ``torch.cat``ing into a single # tensor) so this branch has the *same Python type* as the # unconditioned branch below. ``_merge_data_batches`` picks its # branch from ``values[0]``: when one batch is a Tensor and # another a list it would silently iterate the tensor by dim 0 # via ``for item in v``, producing wrong shapes. Emitting a list # on both paths eliminates the inconsistency at the source. video_tensor = [t.to(device=device) for t in conditioned["video"]] # list of [1,3,T,H,W] sequence_plan = conditioned["sequence_plan"] else: video_tensor = [ torch.zeros(1, 3, sample_args.num_frames, h, w) for _ in range(sample_args.num_outputs) ] # list of [1,3,T,H,W] sequence_plan = None out: dict = { input_vision_key: video_tensor, "image_size": [ torch.tensor([[h, w, h, w]], dtype=torch.float32) for _ in range(sample_args.num_outputs) ], # list of [1,4] "t5_text_embeddings": torch.randn(sample_args.num_outputs, 512, 1024, dtype=torch.bfloat16), # [B,512,1024] "fps": torch.full((sample_args.num_outputs,), float(sample_args.fps)), # [B] "conditioning_fps": torch.full((sample_args.num_outputs,), float(sample_args.fps)), # [B] "num_frames": torch.full((sample_args.num_outputs,), sample_args.num_frames), # [B] "is_preprocessed": True, } if sequence_plan is not None: out["sequence_plan"] = sequence_plan out |= _get_prompt_sample_data(sample_args, model, w=w, h=h, device=device) if prompt_upsampling_image is not None and sequence_plan is not None: out["_prompt_upsampling_images"] = [prompt_upsampling_image.copy() for _ in out[model.input_caption_key]] if sample_args.enable_sound: from cosmos_framework.inference.sound import ( create_placeholder_audio, get_audio_tokenizer_info, inject_sound_into_batch, ) audio_info = get_audio_tokenizer_info(model) if not audio_info.has_sound: raise ValueError("enable_sound=True but model has no sound tokenizer") audio_placeholder = create_placeholder_audio( num_frames=sample_args.num_frames, conditioning_fps=sample_args.fps, audio_info=audio_info, ) inject_sound_into_batch(out, audio_placeholder, model) return out def _merge_data_batches(batches: list[dict[str, Any]]) -> dict[str, Any]: """Merge per-sample data dicts into a single batched dict. Values that are lists are concatenated. Tensors with a batch dimension are concatenated along dim 0. Scalar/bool values are taken from the first batch (and must be equal across all batches). **Aliasing contract.** For a single-batch input, the returned dict is ``batches[0]`` itself (no copy) — this avoids an unnecessary ``torch.cat`` / list rebuild on the hot path. For a multi-batch input, the returned dict is freshly allocated (list values via list-comp, tensor values via ``torch.cat``). The singleton fast path is safe to consume because the in-tree producers already hand this function dicts that the caller fully owns: - ``create_batches_from_dataset`` shallow-copies the per-sample dict once per seed-expanded sibling in ``_expanded_samples``, so the ``batches[0]`` returned here aliases only that sibling's copy — a caller-side top-level mutation cannot leak into sibling samples. - ``_finalize_data_batch`` shallow-copies its input before applying any rename / unbind, so the ``generate_batch`` path is also safe. If you add a new producer that hands shared dict references in, copy at the producer (matching ``_expanded_samples``) rather than removing this fast path — concatenating a length-1 list / single tensor on every singleton batch is meaningful overhead on the hot path. Args: batches (list[dict[str, Any]]): List of data batches to merge. Must be non-empty. Returns: dict[str, Any]: Merged data batch. Raises: ValueError: If the batches have different keys, if values for the same key have inconsistent Python types across batches, or if scalar/bool values are not equal across all batches. """ if len(batches) == 1: return batches[0] # First ensure that all batches have the same keys. reference_keys = set(batches[0].keys()) for i, batch in enumerate(batches[1:], start=1): if set(batch.keys()) != reference_keys: raise ValueError(f"Batch {i} keys {set(batch.keys())} differ from batch 0 keys {reference_keys}") # Then merge the batches. merged: dict[str, Any] = {} keys = batches[0].keys() for key in keys: values = [b[key] for b in batches if key in b] first = values[0] # Without this guard, mixing ``list`` and ``Tensor`` values for the # same key silently produces wrong shapes: the ``isinstance(first, list)`` # branch below uses ``for item in v`` which is also valid on Tensors # (it iterates dim 0), so subsequent Tensor values get unpacked into # the merged list as slices rather than triggering a clear error. if not all(isinstance(v, type(first)) for v in values): raise ValueError( f"Inconsistent value types for key '{key}': " f"{[type(v).__name__ for v in values]}. " "Normalize the type at the source (e.g. always emit a list[Tensor]) " "before calling _merge_data_batches." ) if isinstance(first, list): merged[key] = [item for v in values for item in v] elif isinstance(first, torch.Tensor): if first.ndim <= 0: raise ValueError("Tensor must have at least one (batch) dimension") merged[key] = torch.cat(values, dim=0) else: if not all(v == values[0] for v in values): raise ValueError(f"Key {key} values are not the same: {values}") merged[key] = first return merged def _finalize_sample_args_list(sample_args_list: Sequence[OmniSampleArgs]) -> list[OmniSampleArgs]: """Validate and seed-expand a list of sample args. Behavior is per-sample, so adding samples to (or removing them from) the input list never changes how the *other* samples are handled: - ``num_outputs == 1`` samples are passed through unchanged — the original ``sample_args`` reference (and in particular its ``output_dir``, which may be ``None`` for e.g. padding samples) is preserved. - ``num_outputs > 1`` samples are expanded into ``num_outputs`` fresh deep-copies; each copy gets ``num_outputs = 1``, a unique ``output_dir = original / "{i}"``, and a per-replica seed (``original_seed + i`` if a base seed was provided, else ``None``). Args: sample_args_list: Input samples; may freely mix ``num_outputs == 1`` and ``num_outputs > 1`` entries. Returns: New list of ``OmniSampleArgs`` in input order. Single-output entries are the original references; multi-output entries are fresh deep-copies. Raises: ValueError: If any sample has ``num_outputs > 1`` but no ``output_dir`` (we can't form per-output subdirectories without one). """ finalized_sample_args_list: list[OmniSampleArgs] = [] for sample_args in sample_args_list: if sample_args.num_outputs == 1: finalized_sample_args_list.append(sample_args) continue seed = sample_args.seed num_outputs = sample_args.num_outputs output_dir = sample_args.output_dir if output_dir is None: raise ValueError( f"num_outputs={num_outputs} requires output_dir to be set " f"(sample name={sample_args.name!r}); cannot create " "per-output subdirectories" ) for i in range(num_outputs): sample_args_i = sample_args.model_copy(deep=True) sample_args_i.seed = (seed + i) if seed is not None else None sample_args_i.num_outputs = 1 sample_args_i.output_dir = output_dir / f"{i}" finalized_sample_args_list.append(sample_args_i) return finalized_sample_args_list def create_batches_from_dataset( samples: Iterable[tuple[OmniSampleArgs, dict[str, Any]]], model: OmniMoTModel, *, max_num_seqs: int | None = None, max_model_len: int | None = None, ) -> Generator[tuple[list[OmniSampleArgs], dict[str, Any]]]: """Create batches from pre-loaded (sample_args, data_batch) pairs. Reuses the same token-count / sample-count batching logic as ``OmniInference.create_batches``, but works with dataset iterators that already provide data. Samples with ``num_outputs > 1`` are multi-seed expanded via ``_finalize_sample_args_list``; callers that want only a subset of samples expanded should set ``num_outputs`` accordingly before yielding each sample. Args: samples: Iterable of ``(OmniSampleArgs, data_batch)`` pairs. model: The model, used for token counting and seed expansion. max_num_seqs: Maximum number of sequences per batch. max_model_len: Maximum total tokens per batch. Exactly one of ``max_num_seqs`` or ``max_model_len`` must be set. Yields: ``(sample_args_list, merged_data_batch, per_sample_data_batches)`` tuples. ``per_sample_data_batches`` is the list of individual data dicts before merging, useful when callers need per-sample post-processing. """ # Tensor keys whose non-batch dims may differ across samples and must be # promoted to a length-1 ``list[Tensor]`` so ``_merge_data_batches`` can # flatten them via list-extension instead of failing in ``torch.cat``. # Also, include domain_id so that it is produced as a list[Tensor]. _VARIABLE_SHAPE_TENSOR_KEYS = {"video", "action", "domain_id"} def _prepare_for_merge(db: dict[str, Any]) -> dict[str, Any]: """Reshape a per-sample data dict so ``_merge_data_batches`` can combine it. Returns a shallow copy of ``db`` in which select keys are wrapped in length-1 lists so ``_merge_data_batches`` routes them through its list-concatenation branch (flattening one list per sample into a single batch list) instead of its tensor-``torch.cat`` branch. All other values are passed through by reference. - **``"video"`` / ``"action"``** (see ``_VARIABLE_SHAPE_TENSOR_KEYS``). Input is expected to be a ``[1, *variable_dims]`` tensor — the shape produced by ``_collate_sample``'s ``unsqueeze(0)`` — and is converted to a length-1 ``list[Tensor]`` of shape ``[*variable_dims]``. Without this, samples that share a chunk but have different non-batch dims (e.g. videos at 544x736 vs 640x640, or actions of length 148 vs 104 from variable-length clips in the camera-480 dataset) would fail ``torch.cat``. Downstream consumers (``pack_action``, video tokenization) read these as per-sample lists anyway, so the list form also matches the downstream contract. - **``"domain_id"``**. Each sample carries its own scalar domain id that must remain individually addressable after merging (not concatenated into a single batch tensor) so downstream code can dispatch per sample. Input is expected to be a 0-D tensor and is wrapped in a length-1 ``list[Tensor]`` *preserving the 0-D shape*. Routing a 0-D tensor through ``torch.cat`` directly would otherwise hit ``_merge_data_batches``'s ``ndim <= 0`` guard. Args: db: Single sample's data dict, typically straight out of ``_collate_sample``. Returns: A new dict with the rewrites applied; the original ``db`` is not mutated. Raises: ValueError: If a value at one of the special-cased keys has an unexpected Python type or tensor dimensionality — specifically, a non-``Tensor`` ``video`` / ``action`` / ``domain_id``, or a ``domain_id`` whose ``ndim != 0``. """ updated_db: dict[str, Any] = {} for key, value in db.items(): if key in _VARIABLE_SHAPE_TENSOR_KEYS: if not isinstance(value, torch.Tensor): raise ValueError(f"Expected {key} to be a tensor, got {type(value)}") updated_db[key] = [value.squeeze(0)] else: updated_db[key] = value return updated_db def _expanded_samples() -> Generator[tuple[OmniSampleArgs, dict[str, Any]]]: # Lazily normalize and seed-expand each input sample, yielding # (sample_args, data_batch) pairs ready for budget-batching. # # Each seed-expanded sibling gets its own *shallow copy* of # ``updated_db``. Without this, all siblings of a single source # sample would share the exact same dict reference, and singleton # merge batches (the common case here) would alias straight back # to that dict — so any caller-side top-level key reassignment on # ``merged_batch`` (e.g. ``merged_batch["domain_id"] = [...]``) # would leak into the next sibling. The shallow copy isolates # top-level mutations per sibling at negligible cost; tensor / # list values are still shared between siblings, which is fine # because they're treated as read-only inputs. for sa, db in samples: updated_db = _prepare_for_merge(db) expanded = _finalize_sample_args_list([sa]) for exp_sa in expanded: yield exp_sa, dict(updated_db) for batch in _iter_packed_batches( items=_expanded_samples(), get_sample_args=lambda pair: pair[0], model=model, max_model_len=max_model_len, max_num_seqs=max_num_seqs, ): chunk_args = [pair[0] for pair in batch] chunk_data = [pair[1] for pair in batch] yield chunk_args, _merge_data_batches(chunk_data) def _finalize_data_batch(data_batch: dict[str, Any], batch_size: int, model: OmniMoTModel) -> dict[str, Any]: """Return a finalized + validated copy of *data_batch*. All mutations (key renames, tensor → list unbind) are applied to a fresh shallow copy so the caller's input dict is never modified. This keeps the "no aliasing" responsibility localized at the single place where the mutation happens, instead of forcing every producer of a data dict (e.g. seed-expansion fan-out in ``create_batches_from_dataset``, dummy padding batches in ``create_batches``) to defensively copy before handing the dict to ``generate_batch``. Only the top-level dict structure is copied; tensor / list values inside are shared with the input (which is safe because every mutation here is a top-level key rename or value reassignment, never an in-place op on the value itself). Args: data_batch: Input data dict. Not modified. batch_size: Expected number of samples in the batch (used for validation against the caption-list length). model: Model used to resolve the canonical key names. Returns: New dict with renames + unbinds applied. Raises: ValueError: If both old and new variants of a renamed key are present, or if the post-finalize caption-list length doesn't match ``batch_size``. """ data_batch = dict(data_batch) for old_key, new_key in [ ("video", model.input_video_key), ("images", model.input_image_key), ("ai_caption", model.input_caption_key), ]: if old_key in data_batch and new_key != old_key: if new_key in data_batch: raise ValueError(f"Conflicting keys: '{old_key}' and '{new_key}'") data_batch[new_key] = data_batch.pop(old_key) # Unstack variable length tensors _multi_item_keys = { "text_token_ids", "action", model.input_video_key, model.input_image_key, } for key in _multi_item_keys: if key in data_batch and isinstance(data_batch[key], torch.Tensor): if key == model.input_image_key: data_batch[key] = [ t.unsqueeze(0).squeeze(2) for t in torch.unbind(data_batch[key]) ] # list of [1,C,H,W] elif key == model.input_video_key: if data_batch.get("is_preprocessed", False): data_batch[key] = [t.unsqueeze(0) for t in torch.unbind(data_batch[key])] # list of [1,C,T,H,W] else: data_batch[key] = list(torch.unbind(data_batch[key])) else: data_batch[key] = [[t] for t in torch.unbind(data_batch[key])] # Validate if len(data_batch[model.input_caption_key]) != batch_size: raise ValueError( f"Data batch length ({len(data_batch[model.input_caption_key])}) does not match batch size ({batch_size})" ) return data_batch class SampleDataset(Dataset): """PyTorch map-style dataset over inference sample args. Each item is a ``(SampleArgs, data_dict)`` tuple where the data dict is lazily prepared on access via ``__getitem__``. """ def __init__(self, sample_args_list: Sequence[SampleArgs], model: OmniMoTModel) -> None: self._sample_args_list = list(sample_args_list) self._model = model def __len__(self) -> int: return len(self._sample_args_list) def __getitem__(self, idx: int) -> tuple[SampleArgs, dict[str, Any]]: sample_args = self._sample_args_list[idx] assert isinstance(sample_args, OmniSampleArgs) assert sample_args.output_dir is not None data_batch = sample_args.get_data(device="cuda") if not data_batch: data_batch = get_sample_data(sample_args=sample_args, model=self._model) return sample_args, data_batch @dataclass class OmniInference(Inference): # pyrefly: ignore[bad-override] model: OmniMoTModel vae_decode_stream: torch.cuda.Stream | None = None @property def model_config(self) -> "OmniMoTModelConfig": return self.model.config @classmethod def _get_parallelism_config(cls, setup_args: ParallelismArgs) -> ParallelismConfig: return ParallelismConfig( enable_inference_mode=True, data_parallel_shard_degree=setup_args.dp_shard_size, context_parallel_shard_degree=setup_args.cp_size, cfg_parallel_shard_degree=setup_args.cfgp_size, ) @classmethod def _get_compile_config(cls, setup_args: ParallelismArgs) -> CompileConfig: return CompileConfig( # Translate the flat ``OmniSetupOverrides.use_torch_compile`` public # surface into R1's nested ``CompileConfig.enabled`` knob. enabled=setup_args.use_torch_compile, use_cuda_graphs=setup_args.use_cuda_graphs and setup_args.dp_shard_size * setup_args.cp_size * setup_args.cfgp_size == 1, compiled_region=setup_args.compiled_region, compile_dynamic=setup_args.compile_dynamic, ) @override @classmethod def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self: assert isinstance(setup_args, OmniSetupArgs) assert setup_args.output_dir is not None sampler_override = setup_args.sampler parallelism_config = cls._get_parallelism_config(setup_args) compile_config = cls._get_compile_config(setup_args) if setup_args.checkpoint_type == CheckpointType.DCP and setup_args.config_file_type == ConfigFileType.MODULE: from cosmos_framework.inference.common.config import save_config from cosmos_framework.utils.vfm.model_loader import load_model_from_checkpoint if not setup_args.experiment: raise ValueError("'experiment' is required") if not setup_args.config_file: raise ValueError("'config_file' is required") Cosmos3OmniModel.before_load_model() model, config = load_model_from_checkpoint( experiment_name=setup_args.experiment, config_file=setup_args.config_file, checkpoint_path=setup_args.checkpoint_path, credential_path=setup_args.credential_path or None, parallelism_config=attrs.asdict(parallelism_config), compile_config=attrs.asdict(compile_config), load_ema_to_reg=setup_args.use_ema_weights, experiment_opts=[ *setup_args.experiment_overrides, f"model.config.rectified_flow_inference_config.scheduler_type={sampler_override}", ], use_cache_checkpoint=setup_args.checkpoint_cache_dir is not None, cache_checkpoint_rootdir=str(setup_args.checkpoint_cache_dir or ""), ) model = cast("OmniMoTModel", model) Cosmos3OmniModel.after_load_model(model) save_config(config, setup_args.output_dir) else: checkpoint_path = setup_args.download_checkpoint() if setup_args.config_file_type == ConfigFileType.MODULE: config = None else: model_dict = setup_args.load_model_config_dict() if setup_args.vlm_processor_from_checkpoint: # Source the VLM processor from the loaded checkpoint's own # bundled files instead of the repository hardcoded in the # model config. Drops the redundant base-model download. tokenizer_cfg = model_dict["config"]["vlm_config"]["tokenizer"] tokenizer_cfg.pop("repository", None) tokenizer_cfg.pop("revision", None) tokenizer_cfg.pop("subdir", None) tokenizer_cfg["tokenizer_type"] = str(checkpoint_path) config = Cosmos3OmniConfig(model=model_dict) model = Cosmos3OmniModel.from_pretrained_dcp( checkpoint_path, config=config, parallelism_config=parallelism_config, compile_config=compile_config, ).model if model.config.rectified_flow_inference_config.scheduler_type != sampler_override: model.config.rectified_flow_inference_config.scheduler_type = sampler_override model.set_up_scheduler_and_sampler() log.debug(f"Sampler overridden to: {sampler_override}") vae_decode_stream: torch.cuda.Stream | None = None if setup_args.use_separate_pipeline_vision_decode_gpu: # The CP/CFGP ranks are partitioned into replica-local groups of size # cp_size * cfgp_size. Only the first rank in each group owns separate-VAE # decode work. For example, with cp_size=2 and cfgp_size=1, ranks [0,1] # form one replica and only rank 0 returns True here. replica_size = setup_args.cp_size * setup_args.cfgp_size is_vae_output_rank = (replica_size <= 1) or (get_rank() % replica_size == 0) vae_device_index = setup_args.cp_size * setup_args.cfgp_size if torch.cuda.device_count() <= vae_device_index: raise RuntimeError( "--use-separate-pipeline-vision-decode-gpu requires a spare visible local GPU on the " "same node as the decode-owning rank, but the configured local decode GPU index " f"{vae_device_index} is unavailable with only {torch.cuda.device_count()} visible local GPUs." ) if is_vae_output_rank: vae_device = torch.device("cuda", vae_device_index) inference_device = torch.device("cuda", torch.cuda.current_device()) vae_decode_stream = torch.cuda.Stream(device=vae_device) vae = model.tokenizer_vision_gen.model vae.device = str(vae_device) vae.model = vae.model.to(device=vae_device) vae.scale = tree_map_only(torch.Tensor, lambda tensor: tensor.to(device=vae_device), vae.scale) original_encode = model.encode original_decode = model.decode def encode_on_vae(state: torch.Tensor) -> torch.Tensor: return original_encode(state.to(device=vae_device, non_blocking=True)).to( device=inference_device, non_blocking=True ) def decode_on_vae(latent: torch.Tensor) -> torch.Tensor: return original_decode(latent.to(device=vae_device, non_blocking=True)) model.encode = encode_on_vae model.decode = decode_on_vae log.info( f"Configured vision VAE on device '{vae_device}' while inference remains on '{inference_device}'", rank0_only=False, ) return cls(setup_args=setup_args, model=model, vae_decode_stream=vae_decode_stream, **kwargs) @classmethod def save_data( cls, data: dict[str, Any], *, output_dir: Path, output_name: str, truncate_action_dim: bool = True, ) -> list[Path]: """Save data to disk in multiple formats. Tensors are saved as ``.safetensors``, non-tensor values as ``.pickle``. If ``truncate_action_dim`` is True and both ``action`` and ``raw_action_dim`` are present in ``data``, the action tensor's last dimension is truncated to ``raw_action_dim`` before saving. Returns a list of paths to all files written. """ files: list[Path] = [] data_tensors: dict[str, torch.Tensor] = {} data_pickle: dict[str, Any] = {} for k, v in data.items(): if isinstance(v, list) and len(v) > 0 and isinstance(v[0], torch.Tensor): for i, x in enumerate(v): data_tensors[f"{k}[{i}]"] = x elif isinstance(v, torch.Tensor): data_tensors[k] = v else: data_pickle[k] = v # Truncate `action` tensor's last dimension to `raw_action_dim` if available; # otherwise use the full action tensor as-is. if truncate_action_dim and "action" in data_tensors and "raw_action_dim" in data_tensors: raw_action_dim = data_tensors["raw_action_dim"][0] action = data_tensors["action"][..., :raw_action_dim] data_tensors["action"] = action log.debug(f"Truncated 'action' tensor to shape={action.shape}") if data_tensors: tensors_file = output_dir / f"{output_name}.safetensors" safetensors.torch.save_file( {k: v.detach().cpu().contiguous() for k, v in data_tensors.items()}, tensors_file ) files.append(tensors_file) if data_pickle: pickle_file = output_dir / f"{output_name}.pickle" with pickle_file.open("wb") as f: pickle.dump(data_pickle, f) files.append(pickle_file) return files @override def create_batches( self, sample_args_list: Sequence[SampleArgs] ) -> Generator[tuple[list[SampleArgs], dict[str, Any]]]: assert isinstance(self.setup_args, OmniSetupArgs) max_model_len = self.setup_args.max_model_len max_num_seqs = self.setup_args.max_num_seqs sample_args_list = _finalize_sample_args_list(cast(Sequence[OmniSampleArgs], sample_args_list)) dataset = SampleDataset(sample_args_list, self.model) # Mod-shard the dataset indices across replicas. sampler_indices = list(range(self.replica_id, len(dataset), self.num_replicas)) # --- Phase 1: pre-compute batch boundaries (cheap, no data prep) --- batch_position_lists = list( _iter_packed_batches( items=range(len(sampler_indices)), get_sample_args=lambda pos: sample_args_list[sampler_indices[pos]], model=self.model, max_model_len=max_model_len, max_num_seqs=max_num_seqs, ) ) num_local_batches = len(batch_position_lists) log.debug(f"Number of local batches: {num_local_batches}", rank0_only=False) # --- Phase 2: synchronize batch count across replicas --- # All ranks within a replica share the same replica_id and therefore # the same local batch count, so a global MAX all-reduce is sufficient # to align all replicas. if torch.distributed.is_initialized() and self.num_replicas > 1: count_tensor = torch.tensor([num_local_batches], dtype=torch.long, device="cuda") torch.distributed.all_reduce(count_tensor, op=torch.distributed.ReduceOp.MAX) global_max_batches = int(count_tensor.item()) else: global_max_batches = num_local_batches log.debug(f"Number of global batches: {global_max_batches}") log.debug(f"Number of padding batches: {global_max_batches - num_local_batches}", rank0_only=False) # --- Phase 3: yield real batches (lazily prepare data) --- batches_yielded = 0 for batch_positions in batch_position_lists: chunk_args: list[SampleArgs] = [] chunk_data: list[dict[str, Any]] = [] for pos in batch_positions: sample_idx = sampler_indices[pos] sample_args, data_batch = dataset[sample_idx] if self.setup_args.debug and self.should_process_sample(sample_args): assert sample_args.output_dir is not None sample_args.output_dir.mkdir(parents=True, exist_ok=True) self.save_data( data_batch, output_dir=sample_args.output_dir, output_name="sample_data", ) chunk_args.append(sample_args) chunk_data.append(data_batch) yield chunk_args, _merge_data_batches(chunk_data) batches_yielded += 1 assert batches_yielded == num_local_batches # --- Phase 4: pad with dummy batches so every replica calls # generate_batch the same number of times (prevents collective # deadlocks in context-parallel / CFG-parallel communication). # Minimal-cost padding sample: the dummy batch only exists to keep the # generate_batch call count aligned across replicas, and its output is # discarded (output_dir=None). Force num_steps=1 / guidance=1.0 so it never # raises the per-iteration align_num_steps MAX (which would make the dummy # *and* real samples on peer ranks pad up). The per-step alignment still # pads this dummy up to MAX(real samples), so collective alignment holds; # we just stop inflating that MAX with the (arbitrary) global sample[0]. dummy_sa = sample_args_list[0].model_copy( update={"output_dir": None, "name": "padding", "num_steps": 1, "guidance": 1.0} ) dummy_data = dataset[0][1] while batches_yielded < global_max_batches: yield [dummy_sa], dummy_data batches_yielded += 1 assert batches_yielded == global_max_batches @torch.no_grad() @override def generate_batch( self, sample_args_list: Sequence[SampleArgs], data_batch: dict[str, Any], *, warmup: bool = False ) -> list[SampleOutputs]: assert all(isinstance(sa, OmniSampleArgs) for sa in sample_args_list) transfer_flags = [bool(sa.transfer_hints) for sa in sample_args_list] if any(transfer_flags): assert all(transfer_flags), "Cannot mix transfer and non-transfer samples in a batch" assert len(sample_args_list) == 1, "Batching is not supported for transfer inference" return self._generate_transfer_batch(sample_args_list[0], warmup=warmup) reasoner_flags = [cast(OmniSampleArgs, sa).model_mode.is_reasoner for sa in sample_args_list] if any(reasoner_flags): assert all(reasoner_flags), "Cannot mix reasoner and non-reasoner samples in a batch" return self._generate_reasoner_batch(sample_args_list, data_batch, warmup=warmup) # Process inputs try: with sync_distributed_errors(): for sample_args in sample_args_list: if self.should_process_sample(sample_args) and not warmup: log.debug(f"{sample_args.__class__.__name__}({sample_args})") assert sample_args.output_dir is not None sample_args.output_dir.mkdir(parents=True, exist_ok=True) sample_args_file = sample_args.output_dir / "sample_args.json" sample_args_file.write_text(sample_args.model_dump_json()) log.info(f"Saved sample args to '{sample_args_file}'", rank0_only=False) assert all(sa.num_outputs == 1 for sa in sample_args_list), "num_outputs must be 1" data_batch = _finalize_data_batch( data_batch=data_batch, batch_size=len(sample_args_list), model=self.model ) except Exception as e: return [ self._handle_sample_exception(args, e) for args in sample_args_list if self.should_process_sample(args) and not warmup ] # Generate samples # # Can't catch exceptions here. This code contains collective operations # that will hang if any rank fails. If a rank fails, we must restart # the entire distributed environment. # # Use the first sample's sampling parameters for the whole batch. # All samples in a batch share guidance, num_steps, shift, etc. def _getattr(sample_args_list: Sequence[OmniSampleArgs], attr: str) -> Any: attr_values = [getattr(sa, attr) for sa in sample_args_list] if all(v == attr_values[0] for v in attr_values): return attr_values[0] else: raise ValueError(f"Attribute '{attr}' is not the same for all samples: {attr_values}") is_distilled = self.model.config.fixed_step_sampler_config is not None if is_distilled: sampler = self.model.fixed_step_sampler guidance = 1.0 else: sampler = None guidance = _getattr(sample_args_list, "guidance") should_decode_outputs = self.should_process_sample(sample_args_list[0]) def decode_vision(vision_latent: torch.Tensor) -> torch.Tensor: """ Handles decoding of vision latents, either on the inference device or on a separate VAE device if configured. """ if not should_decode_outputs: tokenizer_vision_gen = self.model.tokenizer_vision_gen return vision_latent.new_zeros( ( vision_latent.shape[0], 3, tokenizer_vision_gen.get_pixel_num_frames(int(vision_latent.shape[2])), int(vision_latent.shape[3]) * tokenizer_vision_gen.spatial_compression_factor, int(vision_latent.shape[4]) * tokenizer_vision_gen.spatial_compression_factor, ) ) if self.vae_decode_stream is None: # We are not using a separate GPU for VAE decoding, so decode directly on the inference device vision = self.model.decode(vision_latent) # [B,C,T,H,W] return ((1.0 + vision) / 2).clamp(0, 1) # [B,C,T,H,W] # We are using a separate GPU for VAE decoding, so we need to issue decode on the VAE device vision_ready = torch.cuda.Event() torch.cuda.current_stream(device=vision_latent.device).record_event(vision_ready) self.vae_decode_stream.wait_event(vision_ready) with torch.cuda.stream(self.vae_decode_stream): vision = self.model.decode(vision_latent) # [B,C,T,H,W] return ((1.0 + vision) / 2).clamp(0, 1) # [B,C,T,H,W] # Use a deterministic fallback (rather than ``random.randint``) when # the caller didn't supply a seed: every rank in a CP / CFG-parallel # replica must compute the same seed for a given sample, otherwise the # initial sampling noise diverges across ranks and the parallel # denoising loop produces corrupt outputs. seed = [sa.seed if sa.seed is not None else _fallback_seed(cast(OmniSampleArgs, sa)) for sa in sample_args_list] outputs: dict[str, Any] | None = None if outputs is None: assert all(sa.num_outputs == 1 for sa in sample_args_list), "num_outputs must be 1" n_sample = sum(cast(OmniSampleArgs, sa).num_outputs for sa in sample_args_list) neg_key = "neg_" + self.model.input_caption_key omni_sample_args_list = cast(Sequence[OmniSampleArgs], sample_args_list) # ``_infer_native_prompt_upsampling_tasks`` owns the full # per-sample decision: opt-in flag, already-upsampled content # check, modality routing, and reasoner capability gating. # It returns a per-sample list of ``UpsampleTask | None``; # ``None`` entries mean "skip native upsampling for this # sample". ``generate_samples_from_batch`` (via # ``_maybe_apply_prompt_upsampling``) consumes the list # directly and dispatches per-task group, so mixed # opted-in / opted-out and mixed-task batches all flow # through without a caller-side collapse. resolved_upsample_tasks = _infer_native_prompt_upsampling_tasks( data_batch, omni_sample_args_list, self.model, ) distinct_upsample_tasks = set(resolved_upsample_tasks) if len(distinct_upsample_tasks) > 1: raise ValueError( "[prompt-upsampling] mixed-task batch: per-sample tasks " f"{resolved_upsample_tasks!r} contain multiple distinct V4.2 " f"tasks {sorted(distinct_upsample_tasks, key=lambda x: x or '')}, but " "`generate_samples_from_batch` carries a single " "`upsample_task` knob. Split the batch by task at the " "caller, or refactor `_maybe_apply_prompt_upsampling` to " "dispatch per-sample." ) upsample_task = next(iter(distinct_upsample_tasks)) # FSDP collective-sequence alignment (throughput-style inference where # ranks hold different samples). Each per-step model forward issues a # param all-gather over the FSDP-shard (dp_shard) group, so if dp_shard # peers disagree on ``num_steps`` that group's collective stream # desyncs and deadlocks NCCL at the watchdog timeout (observed: rank0 # wedged at step 31/50 the instant its dp_shard peer finished 35). # # all_reduce(MAX) the local num_steps over the *dp_shard group* and # pass it as ``align_num_steps``; ranks below the max pad with # discarded dummy steps in generate_samples_from_batch. Scope = the # dp_shard group (not world), because that keeps the reduction within a # single modality: modality must already be homogeneous within any # per-forward collective group (else the forward itself desyncs), and # reasoner-only batches take an early return below and never reach this # collective — a world reduction would deadlock against them. The # per-step CP / CFGP collectives are also covered: cp/cfgp groups # always sit inside one data-parallel replica (replica_id = # rank // (cp*cfgp)), so when dp_shard and the replica block (cp*cfgp) # nest, every cp/cfgp peer lands in a dp_shard group with the same MAX. # The nesting precondition is asserted just below. local_num_steps = _getattr(sample_args_list, "num_steps") align_num_steps = local_num_steps parallel_dims = getattr(self.model, "parallel_dims", None) if ( parallel_dims is not None and parallel_dims.dp_shard_mesh is not None and torch.distributed.is_initialized() and parallel_dims.dp_shard_mesh.size() > 1 ): # Non-nesting CP/CFGP overlays (neither dp_shard nor the cp*cfgp # replica block divides the other) let a cp/cfgp group straddle two # dp_shard groups with different maxima, which a dp_shard-scoped # reduction cannot align. Both presets nest (throughput: cp=cfgp=1; # latency: single replica), so this only guards hand-built layouts. replica_block = parallel_dims.cp * parallel_dims.cfgp dp_shard_sz = parallel_dims.dp_shard if replica_block > 1 and dp_shard_sz % replica_block != 0 and replica_block % dp_shard_sz != 0: raise NotImplementedError( "num_steps collective alignment requires dp_shard " f"({dp_shard_sz}) and cp*cfgp ({replica_block}) to nest " "(one must divide the other). Non-nesting CP/CFGP overlays " "with divergent per-sample num_steps are unsupported." ) _steps_t = torch.tensor( [local_num_steps], device=self.model.tensor_kwargs["device"], dtype=torch.int32 ) torch.distributed.all_reduce( _steps_t, op=torch.distributed.ReduceOp.MAX, group=parallel_dims.dp_shard_mesh.get_group() ) align_num_steps = int(_steps_t.item()) with self._get_timer(f"{self.model.__class__.__name__}.generate_samples_from_batch"): outputs = self.model.generate_samples_from_batch( data_batch, sampler=sampler, guidance=guidance, guidance_interval=_getattr(sample_args_list, "guidance_interval"), seed=seed, num_steps=local_num_steps, align_num_steps=align_num_steps, shift=_getattr(sample_args_list, "shift"), sigma_max=_getattr(sample_args_list, "sigma_max"), has_negative_prompt=neg_key in data_batch, n_sample=n_sample, normalize_cfg=_getattr(sample_args_list, "normalize_cfg"), upsample_task=upsample_task, upsample_max_new_tokens=_getattr(omni_sample_args_list, "prompt_upsampler_max_tokens"), upsample_temperature=_getattr(omni_sample_args_list, "prompt_upsampler_temperature"), upsample_top_k=_getattr(omni_sample_args_list, "prompt_upsampler_top_k"), upsample_top_p=_getattr(omni_sample_args_list, "prompt_upsampler_top_p"), upsample_repetition_penalty=_getattr(omni_sample_args_list, "prompt_upsampler_repetition_penalty"), upsample_presence_penalty=_getattr(omni_sample_args_list, "prompt_upsampler_presence_penalty"), upsample_seed=_getattr(omni_sample_args_list, "prompt_upsampler_seed"), ) with self._get_timer(f"{self.model.__class__.__name__}.decode"): output_vision = outputs.pop("vision") decoded_vision = [decode_vision(vision) for vision in output_vision] outputs["vision"] = [cast(torch.Tensor, vision) for vision in decoded_vision] if self.vae_decode_stream is not None: # If we are using a separate GPU for VAE decoding, wait for results to be ready torch.cuda.current_stream(device=outputs["vision"][0].device).wait_stream(self.vae_decode_stream) for k, v in outputs.items(): if len(v) != len(sample_args_list): raise ValueError(f"Output key '{k}' has length {len(v)} but expected {len(sample_args_list)}") if "sound" in outputs: with self._get_timer(f"{self.model.__class__.__name__}.decode_sound"): outputs["sound"] = [self.model.decode_sound(sound) for sound in outputs.pop("sound")] if warmup: return [] # Save outputs sample_outputs: list[SampleOutputs] = [] try: with sync_distributed_errors(): for sample_idx, sample_args in enumerate(sample_args_list): if self.should_process_sample(sample_args): assert isinstance(sample_args, OmniSampleArgs) assert sample_args.output_dir is not None assert sample_args.num_outputs == 1 output = {k: v[sample_idx].squeeze(0) for k, v in outputs.items()} vision_cthw = output.pop("vision") # Run guardrails self._run_text_guardrail( str(sample_args.output_dir), data_batch[self.model.input_caption_key][sample_idx] ) vision_cthw = self._run_video_guardrail(str(sample_args.output_dir), vision_cthw) output["vision"] = vision_cthw content: dict[str, Any] = {} files: list[Path] = [] # Save debug if self.setup_args.debug: files.extend( self.save_data(output, output_dir=sample_args.output_dir, output_name="output") ) # Save vision if vision_cthw.shape[1] == 1: quality = sample_args.image_save_quality else: quality = sample_args.video_save_quality vision_file = sample_args.output_dir / f"vision{sample_args.vision_extension}" output_fps = sample_args.fps save_img_or_video( vision_cthw, str(vision_file.with_suffix("")), fps=output_fps, quality=quality ) assert vision_file.is_file(), vision_file files.append(vision_file) if "sound" in output: from cosmos_framework.inference.sound import ( get_audio_tokenizer_info, mux_audio_into_video, ) audio_info = get_audio_tokenizer_info(self.model) mux_audio_into_video(vision_file, output["sound"], audio_info.sample_rate) if "action" in output: pred_action = output["action"] if "raw_action_dim" in data_batch: raw_action_dim = int(data_batch["raw_action_dim"][sample_idx].item()) assert pred_action.shape[-1] >= raw_action_dim, ( f"invalid raw_action_dim={raw_action_dim} for action with shape {pred_action.shape}" ) pred_action = pred_action[..., :raw_action_dim] content["action"] = pred_action.detach().cpu().tolist() sample_output = SampleOutputs( args=sample_args.model_dump(mode="json"), outputs=[SampleOutput(content=content, files=files)], ) sample_outputs_file = sample_args.output_dir / "sample_outputs.json" sample_outputs_file.write_text(sample_output.model_dump_json()) log.success(f"Saved sample outputs to '{sample_outputs_file}'", rank0_only=False) sample_outputs.append(sample_output) except Exception as e: return [ self._handle_sample_exception(sample_args, e) for sample_args in sample_args_list if self.should_process_sample(sample_args) ] return sample_outputs @torch.no_grad() def _generate_transfer_batch(self, sample_args: OmniSampleArgs, *, warmup: bool = False) -> list[SampleOutputs]: """Handle transfer inference using the autoregressive generate_transfer_sample path.""" from cosmos_framework.inference.transfer import generate_transfer_sample try: with sync_distributed_errors(): if self.should_process_sample(sample_args) and not warmup: log.debug(f"{sample_args.__class__.__name__}({sample_args})") assert sample_args.output_dir is not None sample_args.output_dir.mkdir(parents=True, exist_ok=True) sample_args_file = sample_args.output_dir / "sample_args.json" sample_args_file.write_text(sample_args.model_dump_json()) log.info(f"Saved sample args to '{sample_args_file}'", rank0_only=False) except Exception as e: if self.should_process_sample(sample_args) and not warmup: return [self._handle_sample_exception(sample_args, e)] return [] transfer_output = generate_transfer_sample(sample_args=sample_args, model=self.model) if warmup: return [] sample_outputs: list[SampleOutputs] = [] try: with sync_distributed_errors(): if self.should_process_sample(sample_args): assert sample_args.output_dir is not None content: dict[str, Any] = {} files: list[Path] = [] vision_cthw = ((1.0 + transfer_output.output_video.squeeze(0)) / 2).clamp(0, 1) if vision_cthw.shape[1] == 1: quality = sample_args.image_save_quality else: quality = sample_args.video_save_quality vision_file = sample_args.output_dir / f"vision{sample_args.vision_extension}" output_fps = transfer_output.fps save_img_or_video(vision_cthw, str(vision_file.with_suffix("")), fps=output_fps, quality=quality) assert vision_file.is_file(), vision_file files.append(vision_file) for hint_key, control_video in transfer_output.control_videos.items(): control_cthw = ((1.0 + control_video.squeeze(0)) / 2).clamp(0, 1) control_file = sample_args.output_dir / f"control_{hint_key}{sample_args.vision_extension}" save_img_or_video( control_cthw, str(control_file.with_suffix("")), fps=output_fps, quality=quality ) files.append(control_file) log.info(f"Saved control video to '{control_file}'", rank0_only=False) sample_output = SampleOutputs( args=sample_args.model_dump(mode="json"), outputs=[SampleOutput(content=content, files=files)], ) sample_outputs_file = sample_args.output_dir / "sample_outputs.json" sample_outputs_file.write_text(sample_output.model_dump_json()) log.success(f"Saved transfer outputs to '{sample_outputs_file}'", rank0_only=False) sample_outputs.append(sample_output) except Exception as e: return [self._handle_sample_exception(sample_args, e)] if self.should_process_sample(sample_args) else [] return sample_outputs @torch.no_grad() def _generate_reasoner_batch( self, sample_args_list: Sequence[SampleArgs], data_batch: dict[str, Any], *, warmup: bool = False, ) -> list[SampleOutputs]: """Reasoner AR text generation. Each prompt writes ``reasoner_text.txt`` and ``SampleOutput.content["reasoner_text"]``. Mixing image-conditioned and text-only samples in one batch is rejected.""" sample_args_list = cast(list[OmniSampleArgs], sample_args_list) prompts: list[str] = data_batch[self.model.input_caption_key] raw_images: list[Image.Image | None] = data_batch["reasoner_images"] n_set = sum(img is not None for img in raw_images) if 0 < n_set < len(raw_images): raise ValueError( "Reasoner batch mixes image-conditioned and text-only samples " f"({n_set}/{len(raw_images)} have vision_path). Split into separate batches." ) images: list[Image.Image] | None = cast(list[Image.Image], raw_images) if n_set == len(raw_images) else None try: with sync_distributed_errors(): for sa, prompt in zip(sample_args_list, prompts): if self.should_process_sample(sa) and not warmup: log.debug(f"{sa.__class__.__name__}({sa})") assert sa.output_dir is not None sa.output_dir.mkdir(parents=True, exist_ok=True) (sa.output_dir / "sample_args.json").write_text(sa.model_dump_json()) self._run_text_guardrail(str(sa.output_dir), prompt) except Exception as e: return [ self._handle_sample_exception(sa, e) for sa in sample_args_list if self.should_process_sample(sa) and not warmup ] # Collective call: every rank must enter so FSDP unshard/reshard and the # cross-rank early-exit reduction stay in lockstep. Not wrapped in try/except. with self._get_timer(f"{self.model.__class__.__name__}.generate_reasoner_text"): texts = self.model.generate_reasoner_text( prompts, max_new_tokens=sample_args_list[0].max_new_tokens, images=images, do_sample=sample_args_list[0].do_sample, temperature=sample_args_list[0].temperature, top_k=sample_args_list[0].top_k, top_p=sample_args_list[0].top_p, repetition_penalty=sample_args_list[0].repetition_penalty, presence_penalty=sample_args_list[0].presence_penalty, seed=sample_args_list[0].seed, ) if warmup: return [] sample_outputs: list[SampleOutputs] = [] try: with sync_distributed_errors(): for sa, text in zip(sample_args_list, texts): if not self.should_process_sample(sa): continue assert sa.output_dir is not None self._run_text_guardrail(str(sa.output_dir), text) txt_path = sa.output_dir / "reasoner_text.txt" txt_path.write_text(text) sample_output = SampleOutputs( args=sa.model_dump(mode="json"), outputs=[SampleOutput(content={"reasoner_text": text}, files=[txt_path])], ) (sa.output_dir / "sample_outputs.json").write_text(sample_output.model_dump_json()) log.success(f"Saved reasoner outputs to '{sa.output_dir}'", rank0_only=False) sample_outputs.append(sample_output) except Exception as e: return [self._handle_sample_exception(sa, e) for sa in sample_args_list if self.should_process_sample(sa)] return sample_outputs @property def replica_size(self) -> int: """ The ranks are divided into computation replicas. The replica size is the product of the context parallelism and CFG parallelism sizes. """ if not hasattr(self.model, "parallel_dims") or self.model.parallel_dims is None: return 1 else: return self.model.parallel_dims.cp_size * self.model.parallel_dims.cfgp_size @property def num_replicas(self) -> int: assert get_world_size() % self.replica_size == 0 return get_world_size() // self.replica_size @property def replica_id(self) -> int: return get_rank() // self.replica_size @property def index_in_replica(self) -> int: return get_rank() % self.replica_size def should_process_sample(self, sample_args: SampleArgs) -> bool: """Whether the sample should be processed by the current rank.""" return sample_args.output_dir is not None and self.index_in_replica == 0 _data_converter = cattrs.preconf.json.make_converter() # torch.Tensor @_data_converter.register_unstructure_hook def _unstructure_torch_tensor(obj: torch.Tensor) -> Any: return { "shape": obj.shape, "dtype": str(obj.dtype), "device": str(obj.device), "values": obj.detach().flatten()[:5].cpu().tolist(), }