nemo_automodel.components.models.muse_spark.vision

NaViT vision tower for Muse Spark Open image and video inputs.

Module Contents

Classes

Name	Description
`MuseSparkVisionAdapter`	Two-layer GELU adapter from pixel-shuffled patches to projector width.
`MuseSparkVisionAttention`	Packed bidirectional attention over independent image/video-frame grids.
`MuseSparkVisionBlock`	One pre-LayerNorm NaViT encoder block.
`MuseSparkVisionEncoder`	NaViT layer stack alternating sparse and full bidirectional attention.
`MuseSparkVisionMLP`	Bias-bearing GELU feed-forward used by NaViT.
`MuseSparkVisionModel`	Variable-resolution Muse Spark Open image/video encoder.
`MuseSparkVisionOutput`	Flattened visual features after pixel shuffle and the adapter.

Functions

Name	Description
`_apply_vision_rope`	Apply adjacent-pair RoPE to Q/K tensors shaped `[tokens,heads,head_dim]`.
`_cumulative_lengths`	Build int32 packed-sequence boundaries for a list of segment lengths.
`_pixel_shuffle_order`	Return block-major indices for one row-major `[H,W]` patch grid.
`_rope_frequencies`	Return adjacent-pair 2D RoPE frequencies shaped `[H*W,D/2,2]`.
`_sparse_order`	Pack one patch grid into fixed spatial blocks for sparse vision layers.

API

class nemo_automodel.components.models.muse_spark.vision.MuseSparkVisionAdapter(
    input_size: int,
    output_size: int,
    dtype: torch.dtype
)

Bases: Module

Two-layer GELU adapter from pixel-shuffled patches to projector width.

c_fc

c_proj

nemo_automodel.components.models.muse_spark.vision.MuseSparkVisionAdapter.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

Map [visual_tokens,input_size] to [visual_tokens,output_size].

class nemo_automodel.components.models.muse_spark.vision.MuseSparkVisionAttention(
    hidden_size: int,
    num_heads: int,
    use_rope: bool,
    dtype: torch.dtype
)

Bases: Module

Packed bidirectional attention over independent image/video-frame grids.

head_dim

= hidden_size // num_heads

k_proj

out_proj

q_proj

scaling

= self.head_dim ** -0.5

v_proj

nemo_automodel.components.models.muse_spark.vision.MuseSparkVisionAttention.forward(
    hidden_states: torch.Tensor,
    frequencies: torch.Tensor,
    cu_seqlens: torch.Tensor
) -> torch.Tensor

Attend over [patches,hidden] using int32 segment boundaries.

class nemo_automodel.components.models.muse_spark.vision.MuseSparkVisionBlock(
    config: dict[str, typing.Any],
    dtype: torch.dtype
)

Bases: GradientCheckpointingLayer

One pre-LayerNorm NaViT encoder block.

attn

ln_1

ln_2

mlp

nemo_automodel.components.models.muse_spark.vision.MuseSparkVisionBlock.forward(
    hidden_states: torch.Tensor,
    frequencies: torch.Tensor,
    cu_seqlens: torch.Tensor
) -> torch.Tensor

Transform packed [patches,hidden] visual states.

class nemo_automodel.components.models.muse_spark.vision.MuseSparkVisionEncoder(
    config: dict[str, typing.Any],
    dtype: torch.dtype
)

Bases: Module

NaViT layer stack alternating sparse and full bidirectional attention.

layer_types

= config.get('layer_types')

layers

sparse_attention_factor

= int(config.get('sparse_attention_factor', 4))

nemo_automodel.components.models.muse_spark.vision.MuseSparkVisionEncoder.forward(
    hidden_states: torch.Tensor,
    frequencies: torch.Tensor,
    full_cu_seqlens: torch.Tensor,
    sparse_cu_seqlens: torch.Tensor
) -> torch.Tensor

Run packed [patches,hidden] states through the NaViT stack.

class nemo_automodel.components.models.muse_spark.vision.MuseSparkVisionMLP(
    hidden_size: int,
    intermediate_size: int,
    dtype: torch.dtype
)

Bases: Module

Bias-bearing GELU feed-forward used by NaViT.

c_fc

c_proj

nemo_automodel.components.models.muse_spark.vision.MuseSparkVisionMLP.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

Transform a [patches,hidden] visual sequence.

class nemo_automodel.components.models.muse_spark.vision.MuseSparkVisionModel(
    config: nemo_automodel.components.models.muse_spark.config.MuseSpark1xOmniConfig
)

Bases: Module

Variable-resolution Muse Spark Open image/video encoder.

adapter

conv1_linear

downsample_factor

= int(navit.get('downsample_factor', 2))

encoder

= MuseSparkVisionEncoder(navit, dtype)

grid_sample_mode

= str(navit.get('grid_sample_mode', 'bilinear'))

head_dim

ln_post

ln_pre

patch_size

= int(navit.get('patch_size', 14))

patch_temporal

= int(navit.get('patch_temporal', 2))

position_grid_h

position_grid_w

positional_embedding_vlm

sparse_attention_factor

= int(navit.get('sparse_attention_factor', 4))

nemo_automodel.components.models.muse_spark.vision.MuseSparkVisionModel.forward(
    pixel_values: torch.Tensor,
    grid_thw: torch.Tensor
) -> nemo_automodel.components.models.muse_spark.vision.MuseSparkVisionOutput

Encode row-major patches.

Parameters:

pixel_values

torch.Tensor

Flattened temporal patches with shape [sum(t*h*w), 3*patch_temporal*patch_size**2].

grid_thw

torch.Tensor

Per-media clip grids with shape [media,3]. Each row is [temporal_clips,grid_height,grid_width].

Returns: MuseSparkVisionOutput

Adapted visual tokens with shape

nemo_automodel.components.models.muse_spark.vision.MuseSparkVisionModel.init_weights(
    buffer_device: torch.device,
    init_std: float
) -> None

Initialize a checkpoint-free vision model on buffer_device.

class nemo_automodel.components.models.muse_spark.vision.MuseSparkVisionOutput(
    last_hidden_state: torch.Tensor
)

Dataclass

Flattened visual features after pixel shuffle and the adapter.

last_hidden_state

Tensor

nemo_automodel.components.models.muse_spark.vision._apply_vision_rope(
    query: torch.Tensor,
    key: torch.Tensor,
    frequencies: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Apply adjacent-pair RoPE to Q/K tensors shaped [tokens,heads,head_dim].

nemo_automodel.components.models.muse_spark.vision._cumulative_lengths(
    lengths: list[int],
    device: torch.device
) -> torch.Tensor

Build int32 packed-sequence boundaries for a list of segment lengths.

nemo_automodel.components.models.muse_spark.vision._pixel_shuffle_order(
    offset: int,
    grid_h: int,
    grid_w: int,
    factor: int,
    device: torch.device
) -> tuple[torch.Tensor, int]

Return block-major indices for one row-major [H,W] patch grid.

nemo_automodel.components.models.muse_spark.vision._rope_frequencies(
    head_dim: int,
    grid_h: int,
    grid_w: int,
    theta: float,
    device: torch.device
) -> torch.Tensor

Return adjacent-pair 2D RoPE frequencies shaped [H*W,D/2,2].

nemo_automodel.components.models.muse_spark.vision._sparse_order(
    offset: int,
    grid_h: int,
    grid_w: int,
    block_h: int,
    block_w: int,
    device: torch.device
) -> tuple[torch.Tensor, list[int], int]

Pack one patch grid into fixed spatial blocks for sparse vision layers.