mstar.model.nemotron_duplex.components.audio_codec

mstar.model.nemotron_duplex.components.audio_codec#

EarTTS RVQ audio codec (tts_model.audio_codec.*).

A DAC/Vocos-style ConvNeXt codec: a downsampling encoder and an upsampling decoder, each a heterogeneous stack of strided resample convs and ConvNeXt blocks, plus a projected-RVQ quantizer (prvq).

Channel schedule (from the checkpoint):

encoder: 18 ->[384]x3 ->768 ->[768]x3 ->1536 ->[1536]x3 ->512(latent) decoder: 512 ->1536 ->[1536]x3 ->768 ->[768]x3 ->384 ->[384]x3 ->18

Forward is a faithful port of NeMo RVQVAEModel / Wav2Latent / Latent2Wav / PreTrainedProbabilisticVQ (nemo/collections/speechlm2/modules/ear_tts_vae_codec.py, branch nemotron-labs-voicechat). Config (from the checkpoint codec_config): n_fft=16 hop_length=4 base_hidden_size=384 channel_mult=(1,2,4) rates=(7,7,9) num_blocks=3 kernel_size=7 latent_size=512 codebook_size=1024 num_quantizers=31 wav_to_token_ratio=1764.

The encoder’s spectrogram/STFT and the decoder’s inverse-STFT overlap-add are ported verbatim. ConvNeXt blocks are causal (left-padded) and use a channel-dim LayerNorm with eps 1e-6.

Classes

AudioCodec([config])

ConvNeXtBlock(dim[, expand, kernel])

Causal 1D ConvNeXt block (ports NeMo ConvNeXt1d).

ProjectedRVQ(num_quantizers, codebook, dim)

Projected residual vector quantizer (ports NeMo PreTrainedProbabilisticVQ).
class mstar.model.nemotron_duplex.components.audio_codec.AudioCodec(config=None)[source]#

Bases: Module

Parameters:

config (CodecConfig | None)

ae_decode(x, constrain_value_range=True)[source]#

Continuous latents [B, T', latent] -> waveform [B, 1, T].

Parameters:
Return type:

Tensor

ae_encode(x)[source]#

Waveform [B, 1, T] -> continuous latents [B, T', latent].

Parameters:

x (Tensor)

Return type:

Tensor

decode(code, code_len=None, constrain_value_range=True)[source]#

Codes [B, T', num_quantizers] -> (waveform [B, 1, T], wav_len).

Parameters:
dequantize(code)[source]#
Parameters:

code (Tensor)

Return type:

Tensor

encode(x, x_len=None)[source]#

Waveform [B, 1, T] -> (codes [B, T', num_quantizers], code_len).

Parameters:
load_weights(weights)[source]#

Load the tts_model.audio_codec.* subtree.

quantize(z)[source]#
Parameters:

z (Tensor)

Return type:

Tensor

static remap(name)[source]#
Parameters:

name (str)

Return type:

str | None

class mstar.model.nemotron_duplex.components.audio_codec.ConvNeXtBlock(dim, expand=4, kernel=7)[source]#

Bases: Module

Causal 1D ConvNeXt block (ports NeMo ConvNeXt1d).

Depthwise conv (left-padded for causality) -> channel-dim LayerNorm -> pointwise expand -> GELU -> pointwise project, with a residual add.

Parameters:
forward(x)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Tensor)

Return type:

Tensor

class mstar.model.nemotron_duplex.components.audio_codec.ProjectedRVQ(num_quantizers, codebook, dim)[source]#

Bases: Module

Projected residual vector quantizer (ports NeMo PreTrainedProbabilisticVQ).

Per-quantizer codebooks (mus_list) and running variances (_variance_list). encode greedily selects the nearest codebook entry per residual level; decode sums the selected embeddings.

Parameters:

  • num_quantizers (int)

  • codebook (int)

  • dim (int)

decode(code)[source]#

[B, T, num_quantizers] codes -> [B, T, dim] continuous latents.

Parameters:

code (Tensor)

Return type:

Tensor

encode(z)[source]#

[B, T, dim] continuous latents -> [B, T, num_quantizers] codes.

Parameters:

z (Tensor)

Return type:

Tensor