nemo_automodel.components.models.muse_spark.audio

Waveform encoder for Muse Spark Open audio tokens.

Module Contents

Classes

Name	Description
`MuseSparkAudioBlock`	Causal convolution followed by a GELU MLP residual branch.
`MuseSparkAudioEncoder`	Encode mono 24-kHz waveforms into one token per 1,920 samples.
`MuseSparkAudioPatchifier`	Patch, project, and normalize one waveform stage.
`MuseSparkAudioRMSNorm`	RMSNorm used by the waveform interaction encoder.
`MuseSparkDepthwiseCausalConv`	Depthwise causal temporal convolution.

Functions

Name	Description
`patchify_audio`	Fold `[batch,samples,channels]` into `[batch,steps,patch*channels]`.

API

class nemo_automodel.components.models.muse_spark.audio.MuseSparkAudioBlock(
    hidden_size: int,
    intermediate_size: int,
    kernel_size: int,
    dtype: torch.dtype
)

Bases: Module

Causal convolution followed by a GELU MLP residual branch.

conv

linear1

linear2

norm

= MuseSparkAudioRMSNorm(hidden_size, dtype)

nemo_automodel.components.models.muse_spark.audio.MuseSparkAudioBlock.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

Transform audio states with shape [batch,time,hidden].

class nemo_automodel.components.models.muse_spark.audio.MuseSparkAudioEncoder(
    output_size: int,
    dtype: torch.dtype,
    patch_size: int = 1920,
    kernel_size: int = 7,
    input_channels: int = 1,
    patchify_sizes: collections.abc.Sequence[int] = (16, 6, 5, 4),
    hidden_sizes: collections.abc.Sequence[int] = (256, 512, 1024, 2048),
    layers_per_stage: collections.abc.Sequence[int] = (2, 2, 2, 2)
)

Bases: Module

Encode mono 24-kHz waveforms into one token per 1,920 samples.

dual_norm

= MuseSparkAudioRMSNorm(output_size, dtype)

norm

= MuseSparkAudioRMSNorm(stage_input, dtype)

o_proj

stages

= nn.Sequential(*stages)

nemo_automodel.components.models.muse_spark.audio.MuseSparkAudioEncoder._encode(
    hidden_states: torch.Tensor
) -> torch.Tensor

Encode a padded batch shaped [batch,samples,channels].

nemo_automodel.components.models.muse_spark.audio.MuseSparkAudioEncoder.forward(
    waveforms: list[torch.Tensor]
) -> list[torch.Tensor]

Encode variable-length mono waveforms into [audio_tokens,output] tensors.

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

Initialize a checkpoint-free audio encoder on buffer_device.

class nemo_automodel.components.models.muse_spark.audio.MuseSparkAudioPatchifier(
    patch_size: int,
    input_size: int,
    output_size: int,
    dtype: torch.dtype
)

Bases: Module

Patch, project, and normalize one waveform stage.

linear

norm

= MuseSparkAudioRMSNorm(output_size, dtype)

nemo_automodel.components.models.muse_spark.audio.MuseSparkAudioPatchifier.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

Transform [batch,time,input] into a shorter projected sequence.

class nemo_automodel.components.models.muse_spark.audio.MuseSparkAudioRMSNorm(
    hidden_size: int,
    dtype: torch.dtype,
    eps: float = 1e-06
)

Bases: Module

RMSNorm used by the waveform interaction encoder.

weight

nemo_automodel.components.models.muse_spark.audio.MuseSparkAudioRMSNorm.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

Normalize audio states shaped [batch,time,hidden] in FP32.

class nemo_automodel.components.models.muse_spark.audio.MuseSparkDepthwiseCausalConv(
    hidden_size: int,
    kernel_size: int,
    dtype: torch.dtype
)

Bases: Module

Depthwise causal temporal convolution.

conv

nemo_automodel.components.models.muse_spark.audio.MuseSparkDepthwiseCausalConv.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

Convolve audio states shaped [batch,time,hidden] causally.

nemo_automodel.components.models.muse_spark.audio.patchify_audio(
    hidden_states: torch.Tensor,
    patch_size: int
) -> torch.Tensor

Fold [batch,samples,channels] into [batch,steps,patch*channels].