# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 # Copied from cosmos3._src.vfm.datasets.sequence_packing import math from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Tuple import torch @dataclass class GenerationDataClean: batch_size: int is_image_batch: bool raw_state_vision: Optional[List[torch.Tensor]] = None x0_tokens_vision: Optional[List[torch.Tensor]] = None fps_vision: Optional[torch.Tensor] = None num_vision_items_per_sample: Optional[List[int]] = None x0_tokens_action: Optional[List[torch.Tensor]] = None fps_action: Optional[torch.Tensor] = None action_domain_id: Optional[List[torch.Tensor]] = None raw_action_dim: Optional[List[Optional[torch.Tensor]]] = None x0_tokens_sound: Optional[List[torch.Tensor]] = None fps_sound: Optional[torch.Tensor] = None # "Fake" types for readability; everything is plain dict at runtime. FactoredSequencePack = dict[str, Any] JointSequencePack = dict[str, Any] SequencePack = FactoredSequencePack | JointSequencePack # ------------------------------------ # Internal helpers # ------------------------------------ def _pad_to_N(N, x: torch.Tensor) -> torch.Tensor: assert x.shape[0] <= N padded = x.new_zeros((N, *x.shape[1:])) padded[: x.shape[0]] = x return padded def _pad( causal_seq: torch.Tensor, full_only_seq: torch.Tensor, max_causal_len: int, max_full_len: int ) -> tuple[torch.Tensor, torch.Tensor]: causal_seq = _pad_to_N(max_causal_len, causal_seq) full_only_seq = _pad_to_N(max_full_len, full_only_seq) return causal_seq, full_only_seq def _compute_mode_indices_and_offsets( split_lens: list[int], attn_modes: list[str], mode: str, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor]: indices = [] offsets = [0] next_offset = 0 start = 0 for split_len, attn_mode in zip(split_lens, attn_modes): if attn_mode == mode: indices.extend(range(start, start + split_len)) next_offset += split_len offsets.append(next_offset) start += split_len return ( torch.tensor(indices, dtype=torch.int32, device=device), torch.tensor(offsets, dtype=torch.int32, device=device), ) def _init_sequence_pack( sample_lens: list[int], split_lens: list[int], attn_modes: list[str], device: torch.device, ) -> dict: _max_sample_len = max(sample_lens) _max_causal_len = max((split_lens[i] for i in range(len(split_lens)) if attn_modes[i] == "causal"), default=0) _max_full_len = max((split_lens[i] for i in range(len(split_lens)) if attn_modes[i] == "full"), default=0) sample_lens_cu = torch.tensor([0] + sample_lens, device=device, dtype=torch.int32) _sample_offsets = torch.cumsum(sample_lens_cu, dim=0, dtype=torch.int32) _causal_indices, _causal_seq_offsets = _compute_mode_indices_and_offsets(split_lens, attn_modes, "causal", device) _full_indices, _full_only_seq_offsets = _compute_mode_indices_and_offsets(split_lens, attn_modes, "full", device) return dict( sample_offsets=_sample_offsets, max_sample_len=_max_sample_len, max_causal_len=_max_causal_len, max_full_len=_max_full_len, _causal_indices=_causal_indices, _full_indices=_full_indices, _causal_seq_offsets=_causal_seq_offsets, _full_only_seq_offsets=_full_only_seq_offsets, _num_causal_tokens=len(_causal_indices), _num_full_tokens=len(_full_indices), split_lens=split_lens, attn_modes=attn_modes, ) def _find_non_causal_text_token_idx( attn_modes: list[str], split_lens: list[int], und_token_indexes: list[int], ) -> list[int]: out = [] full_offset = 0 packed_idx = 0 und_token_set = set(und_token_indexes) for attn_mode, split_len in zip(attn_modes, split_lens): if attn_mode == "full": for local_idx, split_idx in enumerate(range(packed_idx, packed_idx + split_len)): if split_idx in und_token_set: out.append(full_offset + local_idx) full_offset += split_len packed_idx += split_len return out def factored_from_joint_sequence( packed_sequence: torch.Tensor, attn_modes: list[str], split_lens: list[int], sample_lens: list[int], packed_und_token_indexes: torch.Tensor, packed_gen_token_indexes: torch.Tensor, is_image_batch: bool = False, cp_world_size: int = 1, pad_for_cuda_graphs: bool = False, ) -> FactoredSequencePack: non_causal_text_idxs = _find_non_causal_text_token_idx(attn_modes, split_lens, packed_und_token_indexes.tolist()) assert len(non_causal_text_idxs) == 0, "non_causal_text_idxs should be empty" assert sum(sample_lens) == packed_sequence.shape[0] meta = _init_sequence_pack(sample_lens, split_lens, attn_modes, packed_sequence.device) causal_seq = packed_sequence[meta["_causal_indices"]] full_only_seq = packed_sequence[meta["_full_indices"]] return { **meta, "max_num_tokens": sum(sample_lens), "causal_seq": causal_seq, "full_only_seq": full_only_seq, "is_sharded": False, } class SplitInfo: def __init__( self, split_lens: list[int], attn_modes: list[str], sample_lens: list[int], actual_len: int, ): assert sum(sample_lens) == sum(split_lens) max_causal_len = 0 max_full_len = 0 for split_len, attn_mode in zip(split_lens, attn_modes): if attn_mode == "causal": max_causal_len = max(max_causal_len, split_len) elif attn_mode == "full": max_full_len = max(max_full_len, split_len) self.max_causal_len = max_causal_len self.max_full_len = max_full_len self.max_sample_len = max(sample_lens) self.split_lens = split_lens self.attn_modes = attn_modes self.sample_lens = sample_lens def build_packed_sequence( joint_attn_implementation: str, *, packed_sequence: torch.Tensor, attn_modes: list[str], split_lens: list[int], sample_lens: list[int], packed_und_token_indexes: torch.LongTensor, packed_gen_token_indexes: torch.LongTensor, num_heads: int, head_dim: int, num_layers: int, token_shapes=None, natten_parameter_list=None, block_size: int = 128, is_image_batch: bool = False, cp_world_size: int = 1, video_temporal_causal: bool = False, vision_token_shapes=None, action_token_shapes=None, temporal_compression_factor_vision: int = 4, null_action_supertokens: bool = False, pad_for_cuda_graphs: bool = False, ) -> tuple[FactoredSequencePack, SplitInfo, None]: assert joint_attn_implementation == "two_way", ( f"Only two_way attention is supported, got {joint_attn_implementation!r}" ) device = packed_sequence.device attention_meta = SplitInfo( split_lens=split_lens, attn_modes=attn_modes, sample_lens=sample_lens, actual_len=int(packed_sequence.shape[0]), ) input_pack = factored_from_joint_sequence( packed_sequence=packed_sequence, attn_modes=attn_modes, split_lens=split_lens, sample_lens=sample_lens, packed_und_token_indexes=packed_und_token_indexes.to(device), packed_gen_token_indexes=packed_gen_token_indexes.to(device), is_image_batch=is_image_batch, cp_world_size=cp_world_size, pad_for_cuda_graphs=pad_for_cuda_graphs, ) input_pack.pop("split_lens", None) input_pack.pop("attn_modes", None) return input_pack, attention_meta, None def _ensure_core_metadata(pack: SequencePack) -> None: required = [ "sample_offsets", "max_sample_len", "max_causal_len", "max_full_len", "_causal_indices", "_full_indices", "_causal_seq_offsets", "_full_only_seq_offsets", "is_sharded", ] for key in required: if key not in pack: raise KeyError(f"Missing required pack field: {key}") def from_mode_splits( causal_seq: torch.Tensor, full_only_seq: torch.Tensor, orig: FactoredSequencePack | JointSequencePack, is_sharded: bool | None = None, ): """ Create a new sequence pack from two mode splits. Args: causal_seq (torch.Tensor): The causal sequence. full_only_seq (torch.Tensor): The full-only sequence. orig (FactoredSequencePack | JointSequencePack): The metadata source to copy from. is_sharded (bool | None): If True, create a local pack for context parallel. If None, inherits from orig. """ _ensure_core_metadata(orig) if is_sharded is None: is_sharded = orig.get("is_sharded", False) if "packed_sequence" in orig: all_len = int(orig["_causal_indices"].shape[0] + orig["_full_indices"].shape[0]) packed_sequence = causal_seq.new_zeros((all_len, *causal_seq.shape[1:])) # [seq_len,D] packed_sequence[orig["_causal_indices"]] = causal_seq packed_sequence[orig["_full_indices"]] = full_only_seq return from_joint(packed_sequence, orig) else: out = dict(orig) out["causal_seq"] = causal_seq out["full_only_seq"] = full_only_seq out["is_sharded"] = is_sharded return out # ------------------------------------ # Public API # ------------------------------------ def zeros_like(orig: FactoredSequencePack | JointSequencePack, shape: Tuple[int, ...] | torch.Size | None = None): """ Create a new sequence pack with the same metadata as the original, but with all tokens set to zero. Args: orig (FactoredSequencePack | JointSequencePack): The original sequence pack to copy metadata from. shape (Tuple[int, ...] | torch.Size | None): The shape of the new sequence pack. If None, the shape will be the same as the original. """ _ensure_core_metadata(orig) if "packed_sequence" in orig: if shape is None: shape_ = orig["packed_sequence"].shape else: assert len(shape) >= 1 and shape[0] == -1 shape_ = (orig["packed_sequence"].shape[0],) + tuple(shape)[1:] packed_sequence = torch.zeros( shape_, device=orig["packed_sequence"].device, dtype=orig["packed_sequence"].dtype ) # [seq_len,D] return from_joint(packed_sequence, orig) else: if shape is None: shape_causal = orig["causal_seq"].shape shape_full = orig["full_only_seq"].shape else: assert len(shape) >= 1 and shape[0] == -1 shape_causal = (orig["causal_seq"].shape[0],) + tuple(shape)[1:] shape_full = (orig["full_only_seq"].shape[0],) + tuple(shape)[1:] causal_seq = torch.zeros( shape_causal, device=orig["causal_seq"].device, dtype=orig["causal_seq"].dtype ) # [N_causal_tokens,D] full_only_seq = torch.zeros( shape_full, device=orig["full_only_seq"].device, dtype=orig["full_only_seq"].dtype ) # [N_full_tokens,D] return from_mode_splits(causal_seq, full_only_seq, orig) def from_joint(packed_sequence: torch.Tensor, metadata_source: FactoredSequencePack | JointSequencePack): """ Create a new sequence pack from a packed sequence and another sequence pack with the same metadata. Args: packed_sequence (torch.Tensor): Tensor containing all tokens in the batch of sequences. metadata_source (FactoredSequencePack | JointSequencePack): The metadata source to copy from. """ _ensure_core_metadata(metadata_source) if "packed_sequence" in metadata_source: out = dict(metadata_source) out["packed_sequence"] = packed_sequence return out else: if metadata_source["is_sharded"]: # Use sharded sequences as is when is_sharded is True (used in Context Parallel) causal_seq = packed_sequence[: len(metadata_source["causal_seq"])] # [N_causal_tokens,D] full_only_seq = packed_sequence[len(metadata_source["causal_seq"]) :] # [N_full_tokens,D] else: causal_seq = packed_sequence[metadata_source["_causal_indices"]] # [N_causal_tokens,D] full_only_seq = packed_sequence[metadata_source["_full_indices"]] # [N_full_tokens,D] causal_seq, full_only_seq = _pad( causal_seq, full_only_seq, max_causal_len=metadata_source["causal_seq"].shape[0], max_full_len=metadata_source["full_only_seq"].shape[0], ) return from_mode_splits(causal_seq, full_only_seq, metadata_source) def from_und_gen_splits(und_seq: torch.Tensor, gen_seq: torch.Tensor, orig: FactoredSequencePack | JointSequencePack): """ Create a new sequence pack from two und/gen splits. Args: und_seq (torch.Tensor): The understanding sequence. gen_seq (torch.Tensor): The generating sequence. orig (FactoredSequencePack | JointSequencePack): The metadata source to copy from. """ # If we have a joint pack (single packed_sequence), place by und/gen indexes. if "packed_sequence" in orig and "packed_und_token_indexes" in orig and "packed_gen_token_indexes" in orig: all_len = int(und_seq.shape[0] + gen_seq.shape[0]) packed_sequence = und_seq.new_zeros((all_len, *und_seq.shape[1:])) # [seq_len,D] packed_sequence[orig["packed_und_token_indexes"]] = und_seq packed_sequence[orig["packed_gen_token_indexes"]] = gen_seq return from_joint(packed_sequence, orig) # Otherwise, treat und/gen as mode splits (und == causal; gen == full). return from_mode_splits(und_seq, gen_seq, orig) def get_und_seq(pack: SequencePack) -> torch.Tensor: """ Get all understanding tokens in a sequence pack in a single tensor. Args: pack (FactoredSequencePack | JointSequencePack): The sequence pack to get the understanding sequence from. Returns: torch.Tensor: All understanding tokens concatenated over all sequences in the batch. """ if "causal_seq" in pack: return pack["causal_seq"] if "packed_sequence" in pack and "packed_und_token_indexes" in pack: return pack["packed_sequence"][pack["packed_und_token_indexes"]] raise KeyError("Cannot derive und_seq from provided pack") def set_und_seq(pack: SequencePack, value: torch.Tensor) -> None: """ Override the understanding tokens in a sequence pack. The order of tokens passed in must correspond to the order of tokens returned by get_und_seq. Args: pack (FactoredSequencePack | JointSequencePack): The sequence pack to set the understanding sequence in. value (torch.Tensor): The understanding sequence to set. """ if "packed_sequence" in pack and "packed_und_token_indexes" in pack: pack["packed_sequence"][pack["packed_und_token_indexes"]] = value elif "causal_seq" in pack: pack["causal_seq"] = value else: raise KeyError("Cannot set und_seq from provided pack") def get_gen_seq(pack: SequencePack) -> torch.Tensor: """ Get all generating tokens in a sequence pack in a single tensor. Args: pack (FactoredSequencePack | JointSequencePack): The sequence pack to get the generating sequence from. Returns: torch.Tensor: All generating tokens concatenated over all sequences in the batch. """ if "full_only_seq" in pack: return pack["full_only_seq"] if "packed_sequence" in pack and "packed_gen_token_indexes" in pack: return pack["packed_sequence"][pack["packed_gen_token_indexes"]] raise KeyError("Cannot derive gen_seq from provided pack") def set_gen_seq(pack: SequencePack, value: torch.Tensor) -> None: """ Override the generating tokens in a sequence pack. The order of tokens passed in must correspond to the order of tokens returned by get_gen_seq. Args: pack (FactoredSequencePack | JointSequencePack): The sequence pack to set the generating sequence in. value (torch.Tensor): The generating sequence to set. """ if "packed_sequence" in pack and "packed_gen_token_indexes" in pack: pack["packed_sequence"][pack["packed_gen_token_indexes"]] = value elif "full_only_seq" in pack: pack["full_only_seq"] = value else: raise KeyError("Cannot set gen_seq from provided pack") def get_device_and_dtype(pack: SequencePack) -> Tuple[torch.device, torch.dtype]: """ Get the device and dtype of a sequence pack. Args: pack (FactoredSequencePack | JointSequencePack): The sequence pack to get the device and dtype from. Returns: Tuple[torch.device, torch.dtype]: The device and dtype of the sequence pack. """ if "packed_sequence" in pack: return pack["packed_sequence"].device, pack["packed_sequence"].dtype if "causal_seq" in pack and "full_only_seq" in pack: return pack["causal_seq"].device, pack["causal_seq"].dtype raise KeyError("Cannot derive device and dtype from provided pack") def get_all_seq(pack: SequencePack) -> torch.Tensor: """ Get all tokens in a sequence pack in a single tensor. Args: pack (FactoredSequencePack | JointSequencePack): The sequence pack to get the all sequence from. Returns: torch.Tensor: All tokens concatenated over all sequences in the batch. """ if "all_seq" in pack: return pack["all_seq"] if "packed_sequence" in pack: return pack["packed_sequence"] if "causal_seq" in pack and "full_only_seq" in pack: _ensure_core_metadata(pack) if pack["is_sharded"]: assert False, "get_all_seq is not supported in context parallel sharded mode" else: out = pack["causal_seq"].new_zeros( int(pack["_causal_indices"].shape[0] + pack["_full_indices"].shape[0]), *pack["causal_seq"].shape[1:] ) # [seq_len,D] if pack["causal_seq"].shape[0] > 0: out[pack["_causal_indices"]] = pack["causal_seq"][: pack["_causal_indices"].shape[0]] if pack["full_only_seq"].shape[0] > 0: out[pack["_full_indices"]] = pack["full_only_seq"][: pack["_full_indices"].shape[0]] return out raise KeyError("Cannot derive all_seq from provided pack") def get_causal_seq(pack: SequencePack) -> Tuple[torch.Tensor, torch.Tensor]: """ Get the causal sequence and its offsets in a sequence pack. Args: pack (FactoredSequencePack | JointSequencePack): The sequence pack to get the causal sequence from. Returns: Tuple[torch.Tensor, torch.Tensor]: The concatenated causal sub-sequences and the starting offset for each sub-sequence. """ _ensure_core_metadata(pack) if "causal_seq" in pack: return pack["causal_seq"], pack["_causal_seq_offsets"] assert "packed_sequence" in pack return pack["packed_sequence"][pack["_causal_indices"]], pack["_causal_seq_offsets"] def get_full_only_seq(pack: SequencePack) -> Tuple[torch.Tensor, torch.Tensor]: """ Get the full-only sequence and its offsets in a sequence pack. Args: pack (FactoredSequencePack | JointSequencePack): The sequence pack to get the full-only sequence from. Returns: Tuple[torch.Tensor, torch.Tensor]: The concatenated full-only sub-sequences and the starting offset for each sub-sequence. """ _ensure_core_metadata(pack) if "full_only_seq" in pack: return pack["full_only_seq"], pack["_full_only_seq_offsets"] assert "packed_sequence" in pack return pack["packed_sequence"][pack["_full_indices"]], pack["_full_only_seq_offsets"] # ============================================================================ # 3D mRoPE position ID utilities # Copied from cosmos3._src.vfm.models.mot.unified_3dmrope_utils # ============================================================================ def get_3d_mrope_ids_text_tokens( num_tokens: int, temporal_offset: int | float, use_float_positions: bool = False, ) -> tuple[torch.Tensor, int | float]: """Generate 3D mRoPE position IDs for text tokens. For text tokens, all three axes (temporal, height, width) share the same monotonically increasing position IDs, starting from ``temporal_offset``. Args: num_tokens: Number of text tokens. temporal_offset: Current temporal offset to start from. use_float_positions: If True, generate float position IDs. Returns: Tuple of position IDs tensor of shape (3, num_tokens) and updated temporal offset. """ if use_float_positions: ids = torch.arange(num_tokens, dtype=torch.float32) + temporal_offset else: ids = torch.arange(num_tokens, dtype=torch.long) + int(temporal_offset) mrope_ids = ids.unsqueeze(0).expand(3, -1).contiguous() # [3,num_tokens] next_temporal_offset = temporal_offset + num_tokens return mrope_ids, next_temporal_offset def get_3d_mrope_ids_vae_tokens( grid_t: int, grid_h: int, grid_w: int, temporal_offset: int | float, reset_spatial_indices: bool = True, fps: float | None = None, base_fps: float = 24.0, temporal_compression_factor: int = 4, base_temporal_compression_factor: int | None = None, start_frame_offset: int = 0, ) -> tuple[torch.Tensor, int | float]: """Generate 3D mRoPE position IDs for VAE vision tokens (image/video latents). Args: grid_t: Number of temporal frames in the latent grid. grid_h: Height of the latent grid (after patchification). grid_w: Width of the latent grid (after patchification). temporal_offset: Current temporal offset. reset_spatial_indices: If True, spatial indices start from 0 for each vision segment. fps: Frames per second. If None, FPS modulation is disabled. base_fps: Base FPS for normalization. temporal_compression_factor: VAE temporal compression factor. base_temporal_compression_factor: Base temporal compression factor. start_frame_offset: Offset added to frame indices before FPS scaling. Returns: Tuple of position IDs tensor of shape (3, grid_t * grid_h * grid_w) and updated offset. """ fps_modulation_enabled = fps is not None and grid_t > 1 effective_base_tcf = ( base_temporal_compression_factor if base_temporal_compression_factor is not None else temporal_compression_factor ) if fps_modulation_enabled: tps = fps / temporal_compression_factor base_tps = base_fps / effective_base_tcf frame_indices = torch.arange(grid_t, dtype=torch.float32) scaled_t = (frame_indices + start_frame_offset) / tps * base_tps + temporal_offset t_index = scaled_t.view(-1, 1).expand(-1, grid_h * grid_w).flatten() t_dtype = torch.float32 else: t_index = ( torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand(-1, grid_h * grid_w).flatten() + int(temporal_offset) + start_frame_offset ) t_dtype = torch.long h_index = torch.arange(grid_h, dtype=torch.long).view(1, -1, 1).expand(grid_t, -1, grid_w).flatten() w_index = torch.arange(grid_w, dtype=torch.long).view(1, 1, -1).expand(grid_t, grid_h, -1).flatten() if not reset_spatial_indices: spatial_offset = int(temporal_offset) h_index = h_index + spatial_offset w_index = w_index + spatial_offset if fps_modulation_enabled: mrope_ids = torch.stack([t_index, h_index.to(torch.float32), w_index.to(torch.float32)], dim=0) else: mrope_ids = torch.stack([t_index, h_index, w_index], dim=0) max_position = mrope_ids.max().item() next_temporal_offset = math.ceil(max_position) + 1 return mrope_ids, next_temporal_offset # ============================================================================ # Data structures for sequence packing # Copied from cosmos3._src.vfm.datasets.sequence_packing # ============================================================================ @dataclass class ModalityData: """Unified container for a single generation modality's data. This dataclass serves dual purposes: 1. During packing: Acts as a builder, accumulating data in lists 2. After finalize(): Holds finalized tensors ready for model consumption """ sequence_indexes: list[int] | torch.Tensor = field(default_factory=list) timesteps: list[float] | torch.Tensor = field(default_factory=list) mse_loss_indexes: list[int] | torch.Tensor = field(default_factory=list) token_shapes: list = field(default_factory=list) tokens: list[torch.Tensor] = field(default_factory=list) condition_mask: list[torch.Tensor] = field(default_factory=list) noisy_frame_indexes: list[torch.Tensor] = field(default_factory=list) domain_id: list[torch.Tensor] = field(default_factory=list) raw_action_dim: list[torch.Tensor | None] | None = field(default_factory=list) def to_cuda(self) -> None: if isinstance(self.sequence_indexes, torch.Tensor): self.sequence_indexes = self.sequence_indexes.cuda() if isinstance(self.timesteps, torch.Tensor): self.timesteps = self.timesteps.cuda() if isinstance(self.mse_loss_indexes, torch.Tensor): self.mse_loss_indexes = self.mse_loss_indexes.cuda() self.tokens = [token.cuda() for token in self.tokens] self.condition_mask = [cm.cuda() for cm in self.condition_mask] self.noisy_frame_indexes = [ni.cuda() for ni in self.noisy_frame_indexes] self.domain_id = [d.cuda() for d in self.domain_id] if self.raw_action_dim is not None: self.raw_action_dim = [d.cuda() if d is not None else None for d in self.raw_action_dim] @dataclass class PackedSequence: """Unified sequence container - works as builder during packing and final output.""" # Sequence structure sample_lens: list[int] = field(default_factory=list) split_lens: list[int] = field(default_factory=list) attn_modes: list[str] = field(default_factory=list) is_image_batch: bool = False sequence_length: int = 0 # Build-time tracking curr: int = 0 # Text modality (list during build, tensor after finalize) text_ids: list[int] | torch.Tensor = field(default_factory=list) text_indexes: list[int] | torch.Tensor = field(default_factory=list) position_ids: list[int] | torch.Tensor = field(default_factory=list) # Loss computation - Cross Entropy (text) label_ids: list[int] | torch.Tensor | None = field(default_factory=list) ce_loss_indexes: list[int] | torch.Tensor | None = field(default_factory=list) ce_loss_weights: list[float] | torch.Tensor | None = field(default_factory=list) # Build-time mRoPE tracking _use_mrope: bool = False _mrope_temporal_offset: int | float = 0 _mrope_reset_spatial: bool = True # Temporal causal null_action_supertokens: bool = False # Generation modalities vision: ModalityData | None = None action: ModalityData | None = None sound: ModalityData | None = None def finalize( self, gen_data_clean: "GenerationDataClean", ) -> "PackedSequence": """Convert all lists to tensors and compute derived values.""" sequence_length = sum(self.sample_lens) sample_lens = self.sample_lens.copy() split_lens = self.split_lens.copy() attn_modes = self.attn_modes.copy() label_ids: torch.Tensor | None = None ce_loss_indexes: torch.Tensor | None = None ce_loss_weights: torch.Tensor | None = None if self.label_ids and len(self.label_ids) > 0: label_ids = torch.tensor(self.label_ids) ce_loss_indexes = torch.tensor(self.ce_loss_indexes) ce_loss_weights = torch.tensor(self.ce_loss_weights) vision: ModalityData | None = None if self.vision is not None and len(self.vision.sequence_indexes) > 0: vision = ModalityData( sequence_indexes=torch.tensor(self.vision.sequence_indexes, dtype=torch.long), timesteps=torch.tensor(self.vision.timesteps), mse_loss_indexes=torch.tensor(self.vision.mse_loss_indexes, dtype=torch.long), token_shapes=list(self.vision.token_shapes), tokens=self.vision.tokens, condition_mask=list(self.vision.condition_mask), noisy_frame_indexes=list(self.vision.noisy_frame_indexes), ) action: ModalityData | None = None if self.action is not None and len(self.action.sequence_indexes) > 0: action = ModalityData( sequence_indexes=torch.tensor(self.action.sequence_indexes, dtype=torch.long), timesteps=torch.tensor(self.action.timesteps), mse_loss_indexes=torch.tensor(self.action.mse_loss_indexes, dtype=torch.long), token_shapes=list(self.action.token_shapes), tokens=self.action.tokens, condition_mask=list(self.action.condition_mask), noisy_frame_indexes=list(self.action.noisy_frame_indexes), domain_id=( gen_data_clean.action_domain_id if gen_data_clean.action_domain_id is not None else [torch.zeros(1, dtype=torch.long)] * len(self.action.token_shapes) ), raw_action_dim=gen_data_clean.raw_action_dim, ) sound: ModalityData | None = None if self.sound is not None and len(self.sound.sequence_indexes) > 0: sound = ModalityData( sequence_indexes=torch.tensor(self.sound.sequence_indexes, dtype=torch.long), timesteps=torch.tensor(self.sound.timesteps), mse_loss_indexes=torch.tensor(self.sound.mse_loss_indexes, dtype=torch.long), token_shapes=list(self.sound.token_shapes), tokens=self.sound.tokens, condition_mask=list(self.sound.condition_mask), noisy_frame_indexes=list(self.sound.noisy_frame_indexes), ) if self._use_mrope and len(self.position_ids) > 0 and isinstance(self.position_ids[0], torch.Tensor): mrope_tensors: list[torch.Tensor] = self.position_ids # type: ignore[assignment] position_ids = torch.cat(mrope_tensors, dim=1) # [3,actual_seq_len] else: position_ids = torch.tensor(self.position_ids) # [seq_len] return PackedSequence( sequence_length=sequence_length, sample_lens=sample_lens, split_lens=split_lens, attn_modes=attn_modes, is_image_batch=gen_data_clean.is_image_batch, text_ids=torch.tensor(self.text_ids, dtype=torch.long), text_indexes=torch.tensor(self.text_indexes, dtype=torch.long), position_ids=position_ids, label_ids=label_ids, ce_loss_indexes=ce_loss_indexes, ce_loss_weights=ce_loss_weights, vision=vision, action=action, sound=sound, null_action_supertokens=self.null_action_supertokens, ) def to_cuda(self) -> None: if isinstance(self.text_ids, torch.Tensor): self.text_ids = self.text_ids.cuda() if isinstance(self.text_indexes, torch.Tensor): self.text_indexes = self.text_indexes.cuda() if isinstance(self.position_ids, torch.Tensor): self.position_ids = self.position_ids.cuda() if isinstance(self.label_ids, torch.Tensor): self.label_ids = self.label_ids.cuda() if isinstance(self.ce_loss_indexes, torch.Tensor): self.ce_loss_indexes = self.ce_loss_indexes.cuda() if isinstance(self.ce_loss_weights, torch.Tensor): self.ce_loss_weights = self.ce_loss_weights.cuda() if self.vision is not None: self.vision.to_cuda() if self.action is not None: self.action.to_cuda() if self.sound is not None: self.sound.to_cuda() @dataclass class SequencePlan: """Plan describing which modalities are present in a sample.""" has_text: bool has_vision: bool = False condition_frame_indexes_vision: list[int] = field(default_factory=list) has_action: bool = False condition_frame_indexes_action: list[int] = field(default_factory=list) has_sound: bool = False condition_frame_indexes_sound: list[int] = field(default_factory=list) def as_dict(self) -> dict: return { "has_text": self.has_text, "has_vision": self.has_vision, "has_action": self.has_action, "has_sound": self.has_sound, "condition_frame_indexes_vision": self.condition_frame_indexes_vision, "condition_frame_indexes_action": self.condition_frame_indexes_action, "condition_frame_indexes_sound": self.condition_frame_indexes_sound, } # ============================================================================ # Helper functions for packing sequences # ============================================================================ def _pack_text_tokens( packed_seq: PackedSequence, text_ids: List[int], special_tokens: Dict[str, int], curr_rope_id: int, has_generation: bool, use_float_positions: bool = False, ) -> Tuple[int, int, int]: """Pack text tokens into the sequence.""" assert isinstance(packed_seq.text_ids, list), "PackedSequence must be in build mode" assert isinstance(packed_seq.text_indexes, list) assert isinstance(packed_seq.position_ids, list) assert isinstance(packed_seq.label_ids, list) assert isinstance(packed_seq.ce_loss_indexes, list) assert isinstance(packed_seq.ce_loss_weights, list) curr = packed_seq.curr if "bos_token_id" in special_tokens: shifted_text_ids = [special_tokens["bos_token_id"]] + text_ids else: shifted_text_ids = text_ids split_len = 0 packed_seq.text_ids.extend(shifted_text_ids) packed_seq.text_indexes.extend(range(curr, curr + len(shifted_text_ids))) packed_seq.ce_loss_indexes.extend(range(curr, curr + len(shifted_text_ids))) packed_seq.ce_loss_weights.extend([1.0] * len(shifted_text_ids)) packed_seq.label_ids.extend(text_ids[1:] + [special_tokens["eos_token_id"]]) curr += len(shifted_text_ids) split_len += len(shifted_text_ids) packed_seq.text_ids.append(special_tokens["eos_token_id"]) packed_seq.text_indexes.append(curr) curr += 1 split_len += 1 if has_generation: packed_seq.text_ids.append(special_tokens["start_of_generation"]) packed_seq.text_indexes.append(curr) curr += 1 split_len += 1 if packed_seq._use_mrope: text_mrope_ids, packed_seq._mrope_temporal_offset = get_3d_mrope_ids_text_tokens( num_tokens=split_len, temporal_offset=packed_seq._mrope_temporal_offset, use_float_positions=use_float_positions, ) packed_seq.position_ids.append(text_mrope_ids) else: packed_seq.position_ids.extend(range(curr_rope_id, curr_rope_id + split_len)) packed_seq.attn_modes.append("causal") packed_seq.split_lens.append(split_len) packed_seq.curr = curr return curr_rope_id + split_len, split_len, split_len def _pack_vision_tokens( packed_seq: PackedSequence, input_vision_tokens: torch.Tensor, condition_frame_indexes_vision: list[int], input_timestep: float | torch.Tensor, curr_rope_id: int, latent_patch_size: int = 1, vision_fps: float | None = None, enable_fps_modulation: bool = False, base_fps: float = 24.0, temporal_compression_factor: int = 4, ) -> int: """Pack vision tokens into the sequence.""" assert isinstance(packed_seq.position_ids, list), "PackedSequence must be in build mode" curr = packed_seq.curr vision_split_len = 0 if packed_seq.vision is None: packed_seq.vision = ModalityData() assert isinstance(packed_seq.vision.sequence_indexes, list) assert isinstance(packed_seq.vision.mse_loss_indexes, list) assert isinstance(packed_seq.vision.timesteps, list) assert isinstance(packed_seq.vision.tokens, list) _, _, latent_t, latent_h, latent_w = input_vision_tokens.shape if latent_patch_size < 1: raise ValueError(f"latent_patch_size must be >= 1, got {latent_patch_size}") patch_h = math.ceil(latent_h / latent_patch_size) patch_w = math.ceil(latent_w / latent_patch_size) packed_seq.vision.token_shapes.append((latent_t, patch_h, patch_w)) packed_seq.vision.tokens.append(input_vision_tokens) num_vision_tokens = latent_t * patch_h * patch_w packed_seq.vision.sequence_indexes.extend(range(curr, curr + num_vision_tokens)) condition_set = {idx for idx in condition_frame_indexes_vision if 0 <= idx < latent_t} assert isinstance(packed_seq.vision.condition_mask, list) vision_condition_mask = torch.zeros( (latent_t, 1, 1), device=input_vision_tokens.device, dtype=input_vision_tokens.dtype ) for frame_idx in condition_set: vision_condition_mask[frame_idx, 0, 0] = 1.0 packed_seq.vision.condition_mask.append(vision_condition_mask) vision_noisy_frame_indexes = torch.tensor( [idx for idx in range(latent_t) if idx not in condition_set], device=input_vision_tokens.device, dtype=torch.long, ) assert isinstance(packed_seq.vision.noisy_frame_indexes, list) packed_seq.vision.noisy_frame_indexes.append(vision_noisy_frame_indexes) frame_token_stride = patch_h * patch_w for frame_idx in range(latent_t): if frame_idx in condition_set: continue frame_start = curr + frame_idx * frame_token_stride frame_end = frame_start + frame_token_stride packed_seq.vision.mse_loss_indexes.extend(range(frame_start, frame_end)) if isinstance(input_timestep, torch.Tensor): frame_ts = input_timestep[frame_idx].item() else: frame_ts = input_timestep packed_seq.vision.timesteps.extend([frame_ts] * frame_token_stride) curr += num_vision_tokens vision_split_len += num_vision_tokens if packed_seq._use_mrope: effective_fps = vision_fps if enable_fps_modulation else None vision_mrope_ids, packed_seq._mrope_temporal_offset = get_3d_mrope_ids_vae_tokens( grid_t=latent_t, grid_h=patch_h, grid_w=patch_w, temporal_offset=packed_seq._mrope_temporal_offset, reset_spatial_indices=packed_seq._mrope_reset_spatial, fps=effective_fps, base_fps=base_fps, temporal_compression_factor=temporal_compression_factor, ) packed_seq.position_ids.append(vision_mrope_ids) else: packed_seq.position_ids.extend([curr_rope_id] * vision_split_len) packed_seq.curr = curr return vision_split_len def _pack_action_tokens( packed_seq: PackedSequence, input_action_tokens: torch.Tensor, condition_frame_indexes_action: list[int], input_timestep: float, curr_rope_id: int, action_temporal_offset: int | float = 0, enable_fps_modulation: bool = False, base_fps: float = 24.0, action_fps: float | None = None, base_temporal_compression_factor: int | None = None, ) -> int: """Pack action tokens into the sequence.""" assert isinstance(packed_seq.position_ids, list), "PackedSequence must be in build mode" curr = packed_seq.curr action_split_len = input_action_tokens.shape[0] if packed_seq.action is None: packed_seq.action = ModalityData() assert isinstance(packed_seq.action.sequence_indexes, list) assert isinstance(packed_seq.action.mse_loss_indexes, list) assert isinstance(packed_seq.action.timesteps, list) assert isinstance(packed_seq.action.tokens, list) action_indexes = list(range(curr, curr + action_split_len)) packed_seq.action.sequence_indexes.extend(action_indexes) packed_seq.action.token_shapes.append((action_split_len,)) packed_seq.action.tokens.append(input_action_tokens) condition_set = {idx for idx in condition_frame_indexes_action if 0 <= idx < action_split_len} assert isinstance(packed_seq.action.condition_mask, list) action_condition_mask = torch.zeros( (action_split_len, 1), device=input_action_tokens.device, dtype=input_action_tokens.dtype ) for frame_idx in condition_set: action_condition_mask[frame_idx, 0] = 1.0 packed_seq.action.condition_mask.append(action_condition_mask) action_noisy_frame_indexes = torch.tensor( [idx for idx in range(action_split_len) if idx not in condition_set], device=input_action_tokens.device, dtype=torch.long, ) assert isinstance(packed_seq.action.noisy_frame_indexes, list) packed_seq.action.noisy_frame_indexes.append(action_noisy_frame_indexes) frame_token_stride = 1 for frame_idx in range(action_split_len): if frame_idx in condition_set: continue frame_start = curr + frame_idx * frame_token_stride frame_end = frame_start + frame_token_stride packed_seq.action.mse_loss_indexes.extend(range(frame_start, frame_end)) packed_seq.action.timesteps.extend([input_timestep] * frame_token_stride) if packed_seq._use_mrope: effective_fps = action_fps if enable_fps_modulation else None action_mrope_ids, _ = get_3d_mrope_ids_vae_tokens( grid_t=action_split_len, grid_h=1, grid_w=1, temporal_offset=action_temporal_offset, reset_spatial_indices=packed_seq._mrope_reset_spatial, fps=effective_fps, base_fps=base_fps, temporal_compression_factor=1, base_temporal_compression_factor=base_temporal_compression_factor, start_frame_offset=1, ) packed_seq.position_ids.append(action_mrope_ids) else: packed_seq.position_ids.extend([curr_rope_id] * action_split_len) packed_seq.curr = curr + action_split_len return action_split_len def _pack_sound_tokens( packed_seq: PackedSequence, input_sound_tokens: torch.Tensor, condition_frame_indexes_sound: list[int], input_timestep: float, curr_rope_id: int, sound_temporal_offset: int | float = 0, enable_fps_modulation: bool = False, base_fps: float = 24.0, sound_fps: float | None = None, ) -> int: """Pack sound/audio tokens into the sequence.""" assert isinstance(packed_seq.position_ids, list), "PackedSequence must be in build mode" curr = packed_seq.curr _, sound_split_len = input_sound_tokens.shape if packed_seq.sound is None: packed_seq.sound = ModalityData() assert isinstance(packed_seq.sound.sequence_indexes, list) assert isinstance(packed_seq.sound.mse_loss_indexes, list) assert isinstance(packed_seq.sound.timesteps, list) assert isinstance(packed_seq.sound.tokens, list) packed_seq.sound.token_shapes.append((sound_split_len, 1, 1)) packed_seq.sound.sequence_indexes.extend(range(curr, curr + sound_split_len)) packed_seq.sound.tokens.append(input_sound_tokens) condition_set = {idx for idx in condition_frame_indexes_sound if 0 <= idx < sound_split_len} assert isinstance(packed_seq.sound.condition_mask, list) sound_condition_mask = torch.zeros( (sound_split_len, 1), device=input_sound_tokens.device, dtype=input_sound_tokens.dtype ) for frame_idx in condition_set: sound_condition_mask[frame_idx, 0] = 1.0 packed_seq.sound.condition_mask.append(sound_condition_mask) sound_noisy_frame_indexes = torch.tensor( [idx for idx in range(sound_split_len) if idx not in condition_set], device=input_sound_tokens.device, dtype=torch.long, ) assert isinstance(packed_seq.sound.noisy_frame_indexes, list) packed_seq.sound.noisy_frame_indexes.append(sound_noisy_frame_indexes) for frame_idx in range(sound_split_len): if frame_idx in condition_set: continue frame_start = curr + frame_idx frame_end = frame_start + 1 packed_seq.sound.mse_loss_indexes.extend(range(frame_start, frame_end)) packed_seq.sound.timesteps.extend([input_timestep]) if packed_seq._use_mrope: effective_fps = sound_fps if enable_fps_modulation else None sound_mrope_ids, _ = get_3d_mrope_ids_vae_tokens( grid_t=sound_split_len, grid_h=1, grid_w=1, temporal_offset=sound_temporal_offset, reset_spatial_indices=packed_seq._mrope_reset_spatial, fps=effective_fps, base_fps=base_fps, temporal_compression_factor=1, start_frame_offset=0, ) packed_seq.position_ids.append(sound_mrope_ids) else: packed_seq.position_ids.extend([curr_rope_id] * sound_split_len) packed_seq.curr = curr + sound_split_len return sound_split_len def _pack_supertokens_temporal_causal( packed_seq: "PackedSequence", input_vision_tokens: torch.Tensor, input_action_tokens: torch.Tensor | None, condition_frame_indexes_vision: list[int], input_timestep: float | torch.Tensor, curr_rope_id: int, latent_patch_size: int, temporal_compression_factor: int, action_dim: int, vision_fps: float | None = None, action_fps: float | None = None, enable_fps_modulation: bool = False, base_fps: float = 24.0, ) -> tuple[int, bool]: """Pack vision and action tokens in interleaved supertoken order for temporal causal attention. Buffer layout: [action_t0, vision_t0, action_t1, vision_t1, ..., action_{T-1}, vision_{T-1}] """ assert isinstance(packed_seq.position_ids, list), "PackedSequence must be in build mode" _, _, latent_t, latent_h, latent_w = input_vision_tokens.shape patch_h = math.ceil(latent_h / latent_patch_size) patch_w = math.ceil(latent_w / latent_patch_size) tcf = temporal_compression_factor patches_per_frame = patch_h * patch_w supertoken_len = tcf + patches_per_frame if packed_seq.vision is None: packed_seq.vision = ModalityData() if packed_seq.action is None: packed_seq.action = ModalityData() assert isinstance(packed_seq.vision.sequence_indexes, list) assert isinstance(packed_seq.vision.mse_loss_indexes, list) assert isinstance(packed_seq.vision.timesteps, list) assert isinstance(packed_seq.vision.tokens, list) assert isinstance(packed_seq.vision.condition_mask, list) assert isinstance(packed_seq.action.sequence_indexes, list) assert isinstance(packed_seq.action.mse_loss_indexes, list) assert isinstance(packed_seq.action.timesteps, list) assert isinstance(packed_seq.action.tokens, list) assert isinstance(packed_seq.action.condition_mask, list) device = input_vision_tokens.device dtype = input_vision_tokens.dtype null_tokens = torch.zeros(tcf, action_dim, device=device, dtype=dtype) if input_action_tokens is not None: if input_action_tokens.dim() == 3: real_actions = input_action_tokens.squeeze(0) else: real_actions = input_action_tokens if latent_t == 1: all_action_tokens = real_actions else: all_action_tokens = torch.cat([null_tokens, real_actions], dim=0) else: all_action_tokens = torch.zeros(latent_t * tcf, action_dim, device=device, dtype=dtype) null_action_flag = not (latent_t == 1 and input_action_tokens is not None) packed_seq.vision.token_shapes.append((latent_t, patch_h, patch_w)) packed_seq.vision.tokens.append(input_vision_tokens) condition_set_vision = {idx for idx in condition_frame_indexes_vision if 0 <= idx < latent_t} vision_condition_mask = torch.zeros((latent_t, 1, 1), device=device, dtype=dtype) for fidx in condition_set_vision: vision_condition_mask[fidx, 0, 0] = 1.0 packed_seq.vision.condition_mask.append(vision_condition_mask) vision_noisy_frame_indexes = torch.tensor( [idx for idx in range(latent_t) if idx not in condition_set_vision], device=device, dtype=torch.long, ) packed_seq.vision.noisy_frame_indexes.append(vision_noisy_frame_indexes) packed_seq.action.token_shapes.append((latent_t * tcf,)) packed_seq.action.tokens.append(all_action_tokens) action_condition_mask = torch.ones((latent_t * tcf, 1), device=device, dtype=dtype) packed_seq.action.condition_mask.append(action_condition_mask) curr = packed_seq.curr total_split_len = 0 if packed_seq._use_mrope: temporal_offset = packed_seq._mrope_temporal_offset effective_action_fps = action_fps if enable_fps_modulation else None effective_vision_fps = vision_fps if enable_fps_modulation else None fps_active = effective_action_fps is not None t_dtype = torch.float32 if fps_active else torch.long t_offset = float(temporal_offset) if fps_active else int(temporal_offset) null_t = torch.full((tcf,), t_offset, dtype=t_dtype) null_hw = torch.zeros(tcf, dtype=t_dtype) null_ids = torch.stack([null_t, null_hw, null_hw]) # [3,tcf] def _real_action_ids(n_frames: int, start_frame_offset: int) -> torch.Tensor: flat, _ = get_3d_mrope_ids_vae_tokens( grid_t=n_frames * tcf, grid_h=1, grid_w=1, temporal_offset=temporal_offset, reset_spatial_indices=packed_seq._mrope_reset_spatial, fps=effective_action_fps, base_fps=base_fps, temporal_compression_factor=1, base_temporal_compression_factor=tcf, start_frame_offset=start_frame_offset, ) return flat.reshape(3, n_frames, tcf) if latent_t > 1: null_ids_3d = null_ids.reshape(3, 1, tcf) real_ids_3d = _real_action_ids(latent_t - 1, start_frame_offset=1) action_ids_3d = torch.cat([null_ids_3d, real_ids_3d], dim=1) elif input_action_tokens is None: action_ids_3d = null_ids.reshape(3, 1, tcf) else: action_ids_3d = _real_action_ids(1, start_frame_offset=0) vision_ids_flat, new_offset = get_3d_mrope_ids_vae_tokens( grid_t=latent_t, grid_h=patch_h, grid_w=patch_w, temporal_offset=temporal_offset, reset_spatial_indices=packed_seq._mrope_reset_spatial, fps=effective_vision_fps, base_fps=base_fps, temporal_compression_factor=tcf, ) vision_ids_3d = vision_ids_flat.reshape(3, latent_t, patches_per_frame) interleaved_ids = torch.cat([action_ids_3d, vision_ids_3d], dim=2).reshape(3, latent_t * supertoken_len) packed_seq.position_ids.append(interleaved_ids) packed_seq._mrope_temporal_offset = new_offset for frame_t in range(latent_t): action_indexes = list(range(curr, curr + tcf)) packed_seq.action.sequence_indexes.extend(action_indexes) curr += tcf total_split_len += tcf if not packed_seq._use_mrope: packed_seq.position_ids.extend([curr_rope_id] * tcf) frame_indexes = list(range(curr, curr + patches_per_frame)) packed_seq.vision.sequence_indexes.extend(frame_indexes) curr += patches_per_frame total_split_len += patches_per_frame if not packed_seq._use_mrope: packed_seq.position_ids.extend([curr_rope_id] * patches_per_frame) if frame_t not in condition_set_vision: packed_seq.vision.mse_loss_indexes.extend(frame_indexes) frame_ts = input_timestep[frame_t].item() if isinstance(input_timestep, torch.Tensor) else input_timestep packed_seq.vision.timesteps.extend([frame_ts] * patches_per_frame) packed_seq.curr = curr return total_split_len, null_action_flag # ============================================================================ # Main packing functions # ============================================================================ def pack_input_sequence( sequence_plans: list[SequencePlan], input_text_indexes: list[list[int]], gen_data_clean: GenerationDataClean, input_timesteps: torch.Tensor, special_tokens: dict[str, int], max_num_tokens: int | None = None, latent_patch_size: int = 1, skip_text_tokens: bool = False, include_end_of_generation_token: bool = False, position_embedding_type: str = "3d_rope", unified_3d_mrope_reset_spatial_ids: bool = True, unified_3d_mrope_temporal_modality_margin: int = 0, enable_fps_modulation: bool = False, base_fps: float = 24.0, temporal_compression_factor: int = 4, video_temporal_causal: bool = False, action_dim: int = 32, initial_mrope_temporal_offset: int | float = 0, ) -> PackedSequence: """Pack a sequence of input strings and VAE latents into a packed tensor format. Args: sequence_plans: List of SequencePlan items describing which modalities are present. input_text_indexes: List of text token ID sequences. gen_data_clean: GenerationDataClean containing vision, action, and sound tensors. input_timesteps: Diffusion timesteps for each sample. Shape (B,) or (B, 1). special_tokens: Dictionary containing special token IDs. max_num_tokens: Maximum number of tokens (unused, kept for API compatibility). latent_patch_size: Patch size used by the network to pack latents. skip_text_tokens: If True, skip packing text tokens. include_end_of_generation_token: If True, append end-of-generation token. position_embedding_type: Position embedding type for vision tokens. unified_3d_mrope_reset_spatial_ids: If True, spatial indices start from 0 per segment. unified_3d_mrope_temporal_modality_margin: Temporal margin between text and vision. enable_fps_modulation: If True, scale temporal position IDs based on video FPS. base_fps: Base FPS for normalization. temporal_compression_factor: VAE temporal compression factor. video_temporal_causal: If True, pack vision and action as interleaved supertokens. action_dim: Action token dimension for temporal causal packing. initial_mrope_temporal_offset: Initial temporal offset for AR inference. Returns: PackedSequence containing all packed tensors and metadata. """ del max_num_tokens assert special_tokens is not None, "Special tokens must be provided" assert isinstance(input_timesteps, torch.Tensor), "input_timesteps must be a tensor" if input_timesteps.is_cuda: raise ValueError("input_timesteps must be on CPU, not CUDA") if isinstance(input_text_indexes, torch.Tensor): raise ValueError("input_text_tokens must be a list, not a tensor") packed_seq = PackedSequence() packed_seq._use_mrope = position_embedding_type == "unified_3d_mrope" packed_seq._mrope_reset_spatial = unified_3d_mrope_reset_spatial_ids idx_text = 0 idx_vision = 0 idx_action = 0 idx_sound = 0 null_action_flags: list[bool] = [] if not skip_text_tokens: for plan in sequence_plans: assert plan.has_text, "All sequence plans must have has_text=True when skip_text_tokens=False" for sample_idx, sequence_plan in enumerate(sequence_plans): curr_rope_id = 0 sample_len = 0 packed_seq._mrope_temporal_offset = initial_mrope_temporal_offset _ts = input_timesteps[sample_idx] input_timestep = _ts.item() if _ts.numel() == 1 else _ts if sequence_plan.has_text and not skip_text_tokens: text_ids = input_text_indexes[idx_text] idx_text += 1 has_generation_for_sample = sequence_plan.has_vision or sequence_plan.has_action or sequence_plan.has_sound curr_rope_id, _, text_sample_len = _pack_text_tokens( packed_seq, text_ids, special_tokens, curr_rope_id, has_generation=has_generation_for_sample, use_float_positions=enable_fps_modulation, ) sample_len += text_sample_len packed_seq._mrope_temporal_offset += unified_3d_mrope_temporal_modality_margin vision_start_temporal_offset = packed_seq._mrope_temporal_offset if video_temporal_causal and sequence_plan.has_vision: assert position_embedding_type == "unified_3d_mrope", ( "video_temporal_causal=True requires position_embedding_type='unified_3d_mrope'" ) input_vision_tokens = gen_data_clean.x0_tokens_vision[idx_vision] idx_vision += 1 vision_fps = None if ( enable_fps_modulation and gen_data_clean.fps_vision is not None and idx_vision - 1 < len(gen_data_clean.fps_vision) ): vision_fps = float(gen_data_clean.fps_vision[idx_vision - 1].item()) input_action_tokens_tc: torch.Tensor | None = None action_fps_tc: float | None = None if sequence_plan.has_action: input_action_tokens_tc = gen_data_clean.x0_tokens_action[idx_action] if ( enable_fps_modulation and gen_data_clean.fps_action is not None and idx_action < len(gen_data_clean.fps_action) ): action_fps_tc = float(gen_data_clean.fps_action[idx_action].item()) idx_action += 1 supertoken_split_len, null_flag = _pack_supertokens_temporal_causal( packed_seq=packed_seq, input_vision_tokens=input_vision_tokens, input_action_tokens=input_action_tokens_tc, condition_frame_indexes_vision=sequence_plan.condition_frame_indexes_vision, input_timestep=input_timestep, curr_rope_id=curr_rope_id, latent_patch_size=latent_patch_size, temporal_compression_factor=temporal_compression_factor, action_dim=action_dim, vision_fps=vision_fps, action_fps=action_fps_tc, enable_fps_modulation=enable_fps_modulation, base_fps=base_fps, ) null_action_flags.append(null_flag) sample_len += supertoken_split_len vision_split_len = supertoken_split_len action_split_len = 0 else: if sequence_plan.has_vision: num_vis = ( gen_data_clean.num_vision_items_per_sample[sample_idx] if gen_data_clean.num_vision_items_per_sample is not None else 1 ) vision_split_len = 0 for item_idx in range(num_vis): input_vision_tokens = gen_data_clean.x0_tokens_vision[idx_vision] vision_fps: float | None = None if ( enable_fps_modulation and gen_data_clean.fps_vision is not None and idx_vision < len(gen_data_clean.fps_vision) ): vision_fps = float(gen_data_clean.fps_vision[idx_vision].item()) idx_vision += 1 if num_vis > 1 and item_idx < num_vis - 1: latent_t = input_vision_tokens.shape[2] item_condition_frames = list(range(latent_t)) else: item_condition_frames = sequence_plan.condition_frame_indexes_vision item_split_len = _pack_vision_tokens( packed_seq=packed_seq, input_vision_tokens=input_vision_tokens, condition_frame_indexes_vision=item_condition_frames, input_timestep=input_timestep, curr_rope_id=curr_rope_id, latent_patch_size=latent_patch_size, vision_fps=vision_fps, enable_fps_modulation=enable_fps_modulation, base_fps=base_fps, temporal_compression_factor=temporal_compression_factor, ) vision_split_len += item_split_len sample_len += vision_split_len else: vision_split_len = 0 if sequence_plan.has_action: input_action_tokens = gen_data_clean.x0_tokens_action[idx_action] action_fps: float | None = None if ( enable_fps_modulation and gen_data_clean.fps_action is not None and idx_action < len(gen_data_clean.fps_action) ): action_fps = float(gen_data_clean.fps_action[idx_action].item()) idx_action += 1 action_split_len = _pack_action_tokens( packed_seq=packed_seq, input_action_tokens=input_action_tokens, condition_frame_indexes_action=sequence_plan.condition_frame_indexes_action, input_timestep=input_timestep, curr_rope_id=curr_rope_id, action_temporal_offset=vision_start_temporal_offset, enable_fps_modulation=enable_fps_modulation, base_fps=base_fps, action_fps=action_fps, base_temporal_compression_factor=temporal_compression_factor, ) sample_len += action_split_len else: action_split_len = 0 if sequence_plan.has_sound: input_sound_tokens = gen_data_clean.x0_tokens_sound[idx_sound] sound_fps: float | None = None if ( enable_fps_modulation and gen_data_clean.fps_sound is not None and idx_sound < len(gen_data_clean.fps_sound) ): sound_fps = float(gen_data_clean.fps_sound[idx_sound].item()) idx_sound += 1 sound_split_len = _pack_sound_tokens( packed_seq=packed_seq, input_sound_tokens=input_sound_tokens, condition_frame_indexes_sound=sequence_plan.condition_frame_indexes_sound, input_timestep=input_timestep, curr_rope_id=curr_rope_id, sound_temporal_offset=vision_start_temporal_offset, enable_fps_modulation=enable_fps_modulation, base_fps=base_fps, sound_fps=sound_fps, ) sample_len += sound_split_len else: sound_split_len = 0 eov_len = 0 has_any_generation = sequence_plan.has_vision or sequence_plan.has_action or sequence_plan.has_sound if include_end_of_generation_token and has_any_generation: assert isinstance(packed_seq.text_ids, list) assert isinstance(packed_seq.text_indexes, list) assert isinstance(packed_seq.position_ids, list) packed_seq.text_ids.append(special_tokens["end_of_generation"]) packed_seq.text_indexes.append(packed_seq.curr) if packed_seq._use_mrope: eov_dtype = torch.float32 if enable_fps_modulation else torch.long eov_mrope_ids = torch.full((3, 1), packed_seq._mrope_temporal_offset, dtype=eov_dtype) packed_seq.position_ids.append(eov_mrope_ids) # type: ignore[arg-type] packed_seq._mrope_temporal_offset += 1 else: packed_seq.position_ids.append(curr_rope_id) # type: ignore[arg-type] packed_seq.curr += 1 eov_len = 1 sample_len += 1 combined_split_len = vision_split_len + action_split_len + sound_split_len + eov_len packed_seq.attn_modes.append("full") packed_seq.split_lens.append(combined_split_len) packed_seq.sample_lens.append(sample_len) if null_action_flags: assert len(set(null_action_flags)) == 1, ( f"Inconsistent null_action_supertokens across samples: {null_action_flags}." ) packed_seq.null_action_supertokens = null_action_flags[0] return packed_seq.finalize(gen_data_clean=gen_data_clean) def build_sequence_plans_from_data_batch( data_batch: dict, input_video_key, input_image_key: str, ) -> list[SequencePlan]: """Build or retrieve sequence plans from a data batch dictionary. This function extracts sequence plans from the data batch if they exist, otherwise creates default SequencePlan objects for each sample in the batch. Args: data_batch: Dictionary containing the data batch from the dataloader. input_video_key: Key for video tensors in the batch. input_image_key: Key for image tensors in the batch. Returns: List of SequencePlan objects, one per sample in the batch. """ # NOTE: this function is ONLY intended for backward compatibility. # For new modalities, please generate the sequence_plan in the dataset class. if "sequence_plan" in data_batch: return data_batch["sequence_plan"] assert "action" not in data_batch or data_batch["action"] is None, "Action data SHOULD have sequence_plans!" assert "sound" not in data_batch or data_batch["sound"] is None, "Sound data SHOULD have sequence_plans!" batch_size = 0 for key in [input_video_key, input_image_key]: if key in data_batch: val = data_batch[key] if isinstance(val, torch.Tensor): batch_size = val.shape[0] break elif isinstance(val, list): batch_size = len(val) break if batch_size == 0: raise ValueError( f"Cannot determine batch size from data_batch. Expected {input_video_key}, {input_image_key}, or similar key." ) return [ SequencePlan( has_text=True, has_vision=True, condition_frame_indexes_vision=[], ) for _ in range(batch_size) ]