bridge.models.nemotron_omni.nemotron_omni_utils

Module Contents

Functions

patchify_temporal_frame	Resize and normalize one frame for the fixed temporal RADIO policy.
temporal_model_frames	Return frames passed to MCore for one temporal-video sample.
temporal_video_frame_labels	Build one source-timestamped label per temporal tubelet.
processor_patchify_temporal_frames	Apply the public processor’s video resize policy and pack normalized frames.
temporal_tubelet_feature_counts	Compute projected RADIO feature counts in temporal tubelet order.
inference_num_image_tiles	Build image-placeholder replacement counts for pipeline inference.
inference_expanded_image_token_counts	Aggregate projected feature counts for canonical inference prompts.
inference_merged_sequence_length	Return the legacy unpadded length after model-owned vision expansion.
select_inference_next_token	Select the next token from the last real position, excluding PP padding.
load_audio	Load an audio file and resample to target_sr Hz.
_parakeet_feature_extractor	Construct one reusable feature extractor per audio configuration.
valid_audio_feature_lengths	Convert Parakeet feature masks to contiguous-prefix frame lengths.
compute_mel_features_with_length	Convert one waveform to physical mel features and its valid frame length.
compute_mel_features	Convert a raw waveform to a mel spectrogram tensor.
compute_audio_token_count	Compute the expected number of audio tokens for a waveform.

Data

_FrameT
COMPACT_IMAGE_PLACEHOLDER

API

bridge.models.nemotron_omni.nemotron_omni_utils._FrameT

‘TypeVar(…)’

bridge.models.nemotron_omni.nemotron_omni_utils.COMPACT_IMAGE_PLACEHOLDER

‘’

bridge.models.nemotron_omni.nemotron_omni_utils.patchify_temporal_frame(

frame: Any,

*,

height: int,

width: int,

patch_dim: int,

) → torch.Tensor

Resize and normalize one frame for the fixed temporal RADIO policy.

This compatibility helper intentionally places every frame on the supplied canvas. It is shared by fixed-policy training collation and inference so both paths use the same antialiased bicubic interpolation, RADIO normalization, and patch layout.

Parameters:

• frame – PIL-compatible image with convert("RGB") support.

• height – Compatibility-canvas height.

• width – Compatibility-canvas width.

• patch_dim – Vision patch edge length.

Returns:

A tensor with shape [num_patches, 3 * patch_dim * patch_dim].

bridge.models.nemotron_omni.nemotron_omni_utils.temporal_model_frames(

frames: collections.abc.Sequence[bridge.models.nemotron_omni.nemotron_omni_utils._FrameT],

temporal_patch_size: int,

) → list[bridge.models.nemotron_omni.nemotron_omni_utils._FrameT]

Return frames passed to MCore for one temporal-video sample.

MCore pads incomplete groups with the final frame, so odd multi-frame samples retain their original metadata. A single frame needs explicit repetition to select the temporal embedder instead of the image embedder.

Parameters:

• frames – Sampled frames in prompt order.

• temporal_patch_size – Number of frames fused into one temporal tubelet.

Returns:

Frames for model patchification and num_frames metadata.

bridge.models.nemotron_omni.nemotron_omni_utils.temporal_video_frame_labels(

num_frames: int,

*,

temporal_patch_size: int,

source_fps: float | None,

frame_indices: collections.abc.Sequence[int] | None,

) → list[str]

Build one source-timestamped label per temporal tubelet.

Parameters:

• num_frames – Number of sampled frames represented in the prompt.

• temporal_patch_size – Number of consecutive frames per tubelet.

• source_fps – Source-video frame rate, when available.

• frame_indices – Source-frame index for every sampled frame, when available.

Returns:

Frame-label prefixes aligned one-to-one with temporal tubelets.

bridge.models.nemotron_omni.nemotron_omni_utils.processor_patchify_temporal_frames(

frames: collections.abc.Sequence[Any],

*,

image_processor: Any,

patch_dim: int,

) → tuple[torch.Tensor, torch.Tensor]

Apply the public processor’s video resize policy and pack normalized frames.

The Nemotron Omni remote image processor selects a patch-grid size using a separate aspect-preserving video policy. It exposes that policy through the same temporary _is_video_mode switch used by the public multimodal processor. Reusing the processor output here keeps Bridge pixels, grid metadata, and placeholder counts aligned with HF/vLLM preprocessing.

Parameters:

• frames – Frames belonging to one source video, in model order.

• image_processor – Nemotron Omni dynamic-resolution image processor.

• patch_dim – Vision patch edge length.

Returns:

A pair containing packed patches with shape [1, total_patches, 3 * patch_dim**2] and one (height, width) metadata row per frame.

Raises:

ValueError – If the processor contract or returned frame metadata is incompatible with RADIO’s packed temporal input.

bridge.models.nemotron_omni.nemotron_omni_utils.temporal_tubelet_feature_counts(

imgs_sizes: torch.Tensor,

num_frames: torch.Tensor,

*,

temporal_patch_size: int,

patch_dim: int,

pixel_shuffle_factor: int = 2,

) → torch.Tensor

Compute projected RADIO feature counts in temporal tubelet order.

Parameters:

• imgs_sizes – One (height, width) row per ungrouped input frame.

• num_frames – Number of frame rows owned by each image or video item.

• temporal_patch_size – Frames fused into one video tubelet.

• patch_dim – Vision patch edge length.

• pixel_shuffle_factor – Spatial reduction factor per dimension.

Returns:

One feature count per image or temporal tubelet, aligned with compact <img><image></img> wrappers.

Raises:

ValueError – If metadata is inconsistent, a patch grid cannot be shuffled, or frames inside one tubelet use different grids.

bridge.models.nemotron_omni.nemotron_omni_utils.inference_num_image_tiles(

imgs_sizes: torch.Tensor,

*,

patch_dim: int,

pixel_shuffle_factor: int = 2,

num_frames: torch.Tensor | None = None,

temporal_patch_size: int = 1,

) → torch.Tensor

Build image-placeholder replacement counts for pipeline inference.

Dynamic images contribute their post-pixel-shuffle feature count per tile. Temporal tubelets contribute one logical count each; canonical inference applies the fixed tubelet width through

Func:

inference_expanded_image_token_counts. The deprecated LLaVA path instead applies that width inside the model.

Parameters:

• imgs_sizes – Per-image or per-frame (height, width) metadata.

• patch_dim – Vision patch edge length.

• pixel_shuffle_factor – Spatial downsampling factor per dimension.

• num_frames – Frame counts per temporal video, or None for images.

• temporal_patch_size – Frames fused into one temporal tubelet.

Returns:

One integer replacement count per compact image placeholder.

bridge.models.nemotron_omni.nemotron_omni_utils.inference_expanded_image_token_counts(

tile_feature_counts: torch.Tensor,

tiles_per_media: int | collections.abc.Sequence[int] | torch.Tensor,

*,

feature_multiplier: int = 1,

) → torch.Tensor

Aggregate projected feature counts for canonical inference prompts.

Dynamic image processors can split one source image into multiple RADIO tiles, while each <img>...</img> region belongs to the source image. The canonical model needs one <image> placeholder per projected feature, so per-tile counts must be summed back to one count per region. Temporal inference uses one logical tile per tubelet and a fixed feature multiplier for the post-pixel-shuffle tubelet width.

Parameters:

• tile_feature_counts – Number of projected feature rows produced by each RADIO tile or temporal tubelet.

• tiles_per_media – Number of entries in tile_feature_counts owned by each <img>...</img> region.

• feature_multiplier – Additional projected width per count. Use one for dynamic images and the tubelet feature width for temporal video.

Returns:

One expanded placeholder count per <img>...</img> region.

Raises:

ValueError – If counts are non-positive or do not account for every tile/tubelet.

bridge.models.nemotron_omni.nemotron_omni_utils.inference_merged_sequence_length(

input_ids: torch.Tensor,

*,

image_token_index: int,

num_image_tiles: torch.Tensor | None,

image_seq_len: int,

) → int

Return the legacy unpadded length after model-owned vision expansion.

This helper is deprecated because the canonical model consumes an already expanded sequence; its merged length is simply input_ids.shape[1].

Parameters:

• input_ids – One inference prompt row, including generated tokens so far.

• image_token_index – Token ID replaced by vision embeddings.

• num_image_tiles – Row-major replacement metadata per image placeholder.

• image_seq_len – Embeddings contributed by each tile.

Returns:

The real merged sequence length before pipeline padding.

bridge.models.nemotron_omni.nemotron_omni_utils.select_inference_next_token(

logits: torch.Tensor,

merged_sequence_length: int,

) → torch.Tensor

Select the next token from the last real position, excluding PP padding.

bridge.models.nemotron_omni.nemotron_omni_utils.load_audio(path: str, target_sr: int = 16000) → numpy.ndarray

Load an audio file and resample to target_sr Hz.

Supports WAV, MP3, FLAC, and other formats handled by soundfile (with librosa as a fallback for MP3 and other FFmpeg-decoded formats).

Parameters:

• path – Path to the audio file.

• target_sr – Target sampling rate in Hz.

Returns:

1-D float32 numpy array of the mono waveform at target_sr.

bridge.models.nemotron_omni.nemotron_omni_utils._parakeet_feature_extractor(

num_mel_bins: int,

sampling_rate: int,

) → Any

Construct one reusable feature extractor per audio configuration.

bridge.models.nemotron_omni.nemotron_omni_utils.valid_audio_feature_lengths(

attention_mask: torch.Tensor,

*,

num_frames: int,

) → torch.Tensor

Convert Parakeet feature masks to contiguous-prefix frame lengths.

Parakeet retains a padded boundary frame in input_features while its feature-level attention mask records the semantic frame count. Bridge’s sound encoder accepts that mask in compressed form as sound_length.

Parameters:

• attention_mask – Binary mask with shape (batch, frames).

• num_frames – Physical frame width of the corresponding feature tensor.

Returns:

Long tensor containing one valid frame length per batch item.

Raises:

ValueError – If the mask is empty, non-binary, has the wrong shape, or is not a non-empty contiguous prefix for every batch item.

bridge.models.nemotron_omni.nemotron_omni_utils.compute_mel_features_with_length(

waveform: Union[numpy.ndarray, list],

sampling_rate: int = 16000,

num_mel_bins: int = 128,

) → tuple[torch.Tensor, int]

Convert one waveform to physical mel features and its valid frame length.

Parameters:

• waveform – 1-D float32 numpy array (or list) of the mono waveform.

• sampling_rate – Sampling rate of waveform (must match the extractor).

• num_mel_bins – Number of mel frequency bins.

Returns:

A (mel, valid_length) pair. mel has physical shape (frames, num_mel_bins) and may include padded boundary rows; valid_length is derived from Parakeet’s feature attention mask.

bridge.models.nemotron_omni.nemotron_omni_utils.compute_mel_features(

waveform: Union[numpy.ndarray, list],

sampling_rate: int = 16000,

num_mel_bins: int = 128,

) → torch.Tensor

Convert a raw waveform to a mel spectrogram tensor.

Uses HF ParakeetFeatureExtractor (from transformers) to produce mel features compatible with BridgeSoundEncoder / ParakeetEncoder.

Parameters:

• waveform – 1-D float32 numpy array (or list) of the mono waveform.

• sampling_rate – Sampling rate of waveform (must match the extractor).

• num_mel_bins – Number of mel frequency bins.

Returns:

Float tensor of shape (frames, num_mel_bins) – a single clip ready to be batched and passed as sound_clips to the model.

Call :func:compute_mel_features_with_length when the corresponding semantic frame length is also required.

bridge.models.nemotron_omni.nemotron_omni_utils.compute_audio_token_count(

waveform: Union[numpy.ndarray, list],

hop_length: int = 160,

subsampling_factor: int = 8,

) → int

Compute the expected number of audio tokens for a waveform.

Uses the same Conv2D subsampling math as ParakeetEncoder / ParakeetEncoderSubsamplingConv2D: kernel_size=3, stride=2, padding=1, applied log2(subsampling_factor) times to the mel frame count.

Parameters:

• waveform – 1-D waveform array (only its length is used).

• hop_length – Hop length in samples for mel feature extraction.

• subsampling_factor – Subsampling factor of the conformer encoder.

Returns:

Number of audio tokens (at least 1).