mstar.model.cosmos3.components.packing

Joint-sequence packing for Cosmos3 generation (ported from the diffusers Cosmos3OmniPipeline).

Pure, stateless primitives that turn a prompt + latent shape into the transformer’s per-step inputs: the 3D interleaved mRoPE position ids, the text/vision segment layouts, and the chat-template tokenization. Shared by the t2i pipeline and the engine submodule’s input preprocessing. Reproduces the diffusers pipeline’s packed t2i inputs byte-for-byte.

Functions

action_condition_frame_indexes(mode, ...)	Clean (conditioning) action tokens for mode.
action_start_frame_offset(action_length, ...)	mRoPE start-frame offset for the action band: action chunks of length num_frames - 1 start one frame in (aligned to predicted frames 1..N); a full num_frames chunk starts at 0.
build_action_static_inputs(input_ids, ...)	Assemble the static transformer inputs for joint video+action generation.
build_sound_segment(sound_latent_frames, ...)	Sound (AVAE-latent) band of the generation sequence: a (T, 1, 1) grid at the sound tokenizer's latent frame rate, sharing the post-text media temporal offset with the vision band.
build_static_inputs(input_ids, latent_shape, ...)	Assemble the per-prompt static transformer inputs for image/video generation.
build_t2i_static_inputs(input_ids, ...)	Image-mode convenience wrapper around build_static_inputs().
build_text_segment(input_ids, config, device)
build_vision_segment(latent_shape, ...[, ...])	latent_shape is the vision latent tensor shape [B, C, T, H, W].
get_3d_mrope_ids_action_tokens(grid_t, ...)	Action tokens: a frame-rate (T, 1, 1) temporal grid sharing the media offset with the vision band.
get_3d_mrope_ids_text_tokens(num_tokens, ...)	Text tokens: all three axes share the same increasing ids from temporal_offset.
get_3d_mrope_ids_vae_tokens(grid_t, grid_h, ...)	Vision/sound (VAE) tokens: (t, h, w) grid ids, with optional fps modulation of the temporal axis (only when fps is set and grid_t > 1).
normalize_condition_frame_indexes(value, default)	Normalize a request's video-conditioning latent-frame indexes.
tokenize_prompt(tokenizer, prompt, ...[, ...])	Return (cond_input_ids, uncond_input_ids) for image/video generation.
vision_condition_frame_indexes(mode, ...)	Clean (conditioning) latent video frames for mode.

mstar.model.cosmos3.components.packing.action_condition_frame_indexes(mode, action_length)

Clean (conditioning) action tokens for mode.

Parameters:

• mode (str)

• action_length (int)

Return type:

list[int]

mstar.model.cosmos3.components.packing.action_start_frame_offset(action_length, video_length)

mRoPE start-frame offset for the action band: action chunks of length num_frames - 1 start one frame in (aligned to predicted frames 1..N); a full num_frames chunk starts at 0.

Parameters:

• action_length (int)

• video_length (int)

Return type:

int

mstar.model.cosmos3.components.packing.build_action_static_inputs(input_ids, video_latent_shape, action_chunk_size, mode, config, vae_scale_factor_temporal, fps, action_fps, action_start_offset, device)

Assemble the static transformer inputs for joint video+action generation.

The generation sequence is [video tokens | action tokens] after the text prefix. Both media bands share the post-text temporal offset (the 15000 margin), with the action band on its own frame-rate grid. Conditioning per mode decides which video frames and action tokens are clean context vs noisy/predicted (see vision_condition_frame_indexes() / action_condition_frame_indexes()).

Parameters:

• input_ids (list[int])

• video_latent_shape (tuple[int, int, int, int, int])

• action_chunk_size (int)

• mode (str)

• vae_scale_factor_temporal (int)

• fps (float)

• action_fps (float)

• action_start_offset (int)

Return type:

dict[str, Any]

mstar.model.cosmos3.components.packing.build_sound_segment(sound_latent_frames, mrope_offset, curr, config, device)

Sound (AVAE-latent) band of the generation sequence: a (T, 1, 1) grid at the sound tokenizer’s latent frame rate, sharing the post-text media temporal offset with the vision band. Every sound frame is noisy/predicted (sound generation has no clean conditioning tokens).

Parameters:

• sound_latent_frames (int)

• mrope_offset (int | float)

• curr (int)

Return type:

dict[str, Any]

mstar.model.cosmos3.components.packing.build_static_inputs(input_ids, latent_shape, config, vae_scale_factor_temporal, fps, device, has_image_condition=False, sound_latent_frames=None, noisy_frames=None)

Assemble the per-prompt static transformer inputs for image/video generation. latent_shape is [B, C, T, H, W] (T == 1 for images; T == 1 + (num_frames - 1) // temporal_factor for video). When has_image_condition is set, latent frame 0 is a clean conditioning anchor (image-to-video): it stays in the sequence but is excluded from the noisy / predicted frames. noisy_frames instead lists the predicted latent frames explicitly (video-to-video, where the conditioned frames are the complement). sound_latent_frames appends a jointly denoised sound band after the vision tokens (the generation block becomes [vision | sound]). Step-varying fields (vision_tokens, vision_timesteps) are spliced in per denoising step by the caller.

Parameters:

• input_ids (list[int])

• latent_shape (tuple[int, int, int, int, int])

• vae_scale_factor_temporal (int)

• fps (float)

• has_image_condition (bool)

• sound_latent_frames (int | None)

• noisy_frames (list[int] | None)

Return type:

dict[str, Any]

mstar.model.cosmos3.components.packing.build_t2i_static_inputs(input_ids, latent_shape, config, vae_scale_factor_temporal, fps, device)

Image-mode convenience wrapper around build_static_inputs().

Parameters:

• input_ids (list[int])

• latent_shape (tuple[int, int, int, int, int])

• vae_scale_factor_temporal (int)

• fps (float)

Return type:

dict[str, Any]

mstar.model.cosmos3.components.packing.build_text_segment(input_ids, config, device)

Parameters:

input_ids (list[int])

Return type:

dict[str, Any]

mstar.model.cosmos3.components.packing.build_vision_segment(latent_shape, has_image_condition, mrope_offset, vision_fps, curr, config, vae_scale_factor_temporal, device, noisy_frames=None)

latent_shape is the vision latent tensor shape [B, C, T, H, W].

noisy_frames lists the latent frames that are noisy (predicted); the rest are clean conditioning context. When None it defaults to frame 0 clean if has_image_condition else all frames noisy — i.e. the t2i/t2v/i2v layouts. Action modes pass an explicit list (e.g. [] for inverse-dynamics, where the whole video is conditioning).

Parameters:

• latent_shape (tuple[int, int, int, int, int])

• has_image_condition (bool)

• mrope_offset (int | float)

• vision_fps (float | None)

• curr (int)

• vae_scale_factor_temporal (int)

• noisy_frames (list[int] | None)

Return type:

dict[str, Any]

mstar.model.cosmos3.components.packing.get_3d_mrope_ids_action_tokens(grid_t, temporal_offset, action_fps, base_fps=24.0, base_temporal_compression_factor=4, start_frame_offset=1)

Action tokens: a frame-rate (T, 1, 1) temporal grid sharing the media offset with the vision band. The action stream is uncompressed in time (temporal_compression_factor=1) but its rate is expressed in the same base-fps units as the vision latents (base_temporal_compression_factor), so an action chunk lines up temporally with the conditioning video.

Parameters:

• grid_t (int)

• temporal_offset (int | float)

• action_fps (float | None)

• base_fps (float)

• base_temporal_compression_factor (int)

• start_frame_offset (int)

Return type:

tuple[Tensor, int | float]

mstar.model.cosmos3.components.packing.get_3d_mrope_ids_text_tokens(num_tokens, temporal_offset, use_float_positions=False)

Text tokens: all three axes share the same increasing ids from temporal_offset.

Parameters:

• num_tokens (int)

• temporal_offset (int | float)

• use_float_positions (bool)

Return type:

tuple[Tensor, int | float]

mstar.model.cosmos3.components.packing.get_3d_mrope_ids_vae_tokens(grid_t, grid_h, grid_w, temporal_offset, reset_spatial_indices=True, fps=None, base_fps=24.0, temporal_compression_factor=4, base_temporal_compression_factor=None, start_frame_offset=0)

Vision/sound (VAE) tokens: (t, h, w) grid ids, with optional fps modulation of the temporal axis (only when fps is set and grid_t > 1).

Parameters:

• grid_t (int)

• grid_h (int)

• grid_w (int)

• temporal_offset (int | float)

• reset_spatial_indices (bool)

• fps (float | None)

• base_fps (float)

• temporal_compression_factor (int)

• base_temporal_compression_factor (int | None)

• start_frame_offset (int)

Return type:

tuple[Tensor, int | float]

mstar.model.cosmos3.components.packing.normalize_condition_frame_indexes(value, default)

Normalize a request’s video-conditioning latent-frame indexes.

Accepts an int, a comma-separated string, or an iterable of ints; returns a sorted, de-duplicated tuple. None falls back to default.

Parameters:

default (tuple[int, ...])

Return type:

tuple[int, …]

mstar.model.cosmos3.components.packing.tokenize_prompt(tokenizer, prompt, negative_prompt, num_frames, height, width, fps=24.0, use_system_prompt=True, add_resolution_template=True, add_duration_template=True, max_sequence_length=None)

Return (cond_input_ids, uncond_input_ids) for image/video generation.

Mirrors the diffusers pipeline: apply the Qwen2 chat template with the mode-specific system prompt and metadata sentences (duration for video, then resolution), using inverse templates for the negative prompt, then append the eos + start-of-generation (<|vision_start|>) special tokens. Image mode is num_frames == 1. max_sequence_length caps the chat-template token count BEFORE the two trailing markers (the reference serving pipeline’s truncation contract); None leaves the prompt unbounded.

Parameters:

• prompt (str)

• negative_prompt (str | None)

• num_frames (int)

• height (int)

• width (int)

• fps (float)

• use_system_prompt (bool)

• add_resolution_template (bool)

• add_duration_template (bool)

• max_sequence_length (int | None)

Return type:

tuple[list[int], list[int]]

mstar.model.cosmos3.components.packing.vision_condition_frame_indexes(mode, latent_frames)

Clean (conditioning) latent video frames for mode.

Parameters:

• mode (str)

• latent_frames (int)

Return type:

list[int]