bridge.models.nemotron_omni.nemotron_omni_utils

Module Contents

Functions

patchify_temporal_frame	Resize and normalize one frame for MCore’s square temporal RADIO path.
temporal_model_frames	Return frames passed to MCore for one temporal-video sample.
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 MCore’s square temporal RADIO path.

The public HF processor preserves aspect ratio, but pinned MCore requires temporal tubelets to share one square spatial grid. This helper is shared by 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.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).