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_merged_sequence_length

Return the unpadded sequence length after vision-token replacement.

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.

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.

The first pipeline stage can derive these counts from vision encoder outputs, but the last stage needs the same row-major metadata to expand input positions. Dynamic images contribute their post-pixel-shuffle token count per compact placeholder. Temporal tubelets contribute one tile each; LLaVAModel.img_seq_len supplies their fixed embedding width.

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_merged_sequence_length(
input_ids: torch.Tensor,
*,
image_token_index: int,
num_image_tiles: torch.Tensor | None,
image_seq_len: int,
) → int#

Return the unpadded sequence length after vision-token replacement.

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.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.

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).