mstar.model.voxcpm2.components.audio_vae

AudioVAE V2 of VoxCPM2: 16 kHz waveform -> 25 Hz latents -> 48 kHz waveform.

A port of the reference AudioVAEV2 (a DAC-style causal conv VAE with an asymmetric decoder and sample-rate conditioning). Every convolution pads on the left only and every transposed convolution trims on the right, so a latent frame’s audio depends on a bounded span of earlier frames: decoder_receptive_field_frames() derives that span from the modules, which is how much left context a streamed window needs to reproduce a whole decode.

Weight norm is folded into plain weights at load (fold_weight_norm); the checkpoint’s weight_g / weight_v pairs are combined once, in fp32.

Functions

decoder_receptive_field_frames(decoder)	Latent frames of left context a frame's audio needs to equal a whole decode.
fold_weight_norm(state)	weight_g / weight_v pairs -> weight (torch.nn.utils.weight_norm, dim 0).

Classes

AudioVAE(config)	encode: [B, samples@16k] -> [B, 64, frames]; decode: [B, 64, frames] -> [B, frames * 1920] at 48 kHz.
CausalConv1d(*args[, padding, output_padding])
CausalDecoder(config)
CausalDecoderBlock(input_dim, output_dim, stride)
CausalEncoder(d_model, latent_dim, strides, ...)
CausalEncoderBlock(output_dim, stride[, groups])
CausalResidualUnit(dim[, dilation, kernel, ...])
CausalTransposeConv1d(*args[, padding, ...])
SampleRateCondition(input_dim, buckets)	scale_bias conditioning: a per-channel affine picked by output-rate bucket.
Snake1d(channels)

class mstar.model.voxcpm2.components.audio_vae.AudioVAE(config)

Bases: Module

encode: [B, samples@16k] -> [B, 64, frames]; decode: [B, 64, frames] -> [B, frames * 1920] at 48 kHz.

Parameters:

config (AudioVAEConfig)

decode(z, sr_index=None)

Parameters:

• z (Tensor)

• sr_index (Tensor | None)

Return type:

Tensor

encode(audio)

Parameters:

audio (Tensor)

Return type:

Tensor

output_rate_index(device)

Bucket of the default output rate, as AudioVAEV2.decode picks it.

Return type:

Tensor

class mstar.model.voxcpm2.components.audio_vae.CausalConv1d(*args, padding=0, output_padding=0, **kwargs)

Bases: Conv1d

Parameters:

• padding (str | tuple[int, ...])

• output_padding (tuple[int, ...])

forward(x)

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.voxcpm2.components.audio_vae.CausalDecoder(config)

Bases: Module

Parameters:

config (AudioVAEConfig)

forward(z, sr_index)

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:

• z (Tensor)

• sr_index (Tensor)

Return type:

Tensor

class mstar.model.voxcpm2.components.audio_vae.CausalDecoderBlock(input_dim, output_dim, stride, groups=1)

Bases: Module

Parameters:

• input_dim (int)

• output_dim (int)

• stride (int)

• groups (int)

forward(x)

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.voxcpm2.components.audio_vae.CausalEncoder(d_model, latent_dim, strides, depthwise)

Bases: Module

Parameters:

• d_model (int)

• latent_dim (int)

• strides (tuple[int, ...])

• depthwise (bool)

forward(x)

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.voxcpm2.components.audio_vae.CausalEncoderBlock(output_dim, stride, groups=1)

Bases: Module

Parameters:

• output_dim (int)

• stride (int)

• groups (int)

forward(x)

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.voxcpm2.components.audio_vae.CausalResidualUnit(dim, dilation=1, kernel=7, groups=1)

Bases: Module

Parameters:

• dim (int)

• dilation (int)

• kernel (int)

• groups (int)

forward(x)

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.voxcpm2.components.audio_vae.CausalTransposeConv1d(*args, padding=0, output_padding=0, **kwargs)

Bases: ConvTranspose1d

Parameters:

• padding (str | tuple[int, ...])

• output_padding (tuple[int, ...])

forward(x)

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.voxcpm2.components.audio_vae.SampleRateCondition(input_dim, buckets)

Bases: Module

scale_bias conditioning: a per-channel affine picked by output-rate bucket.

Parameters:

• input_dim (int)

• buckets (int)

forward(x, sr_index)

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)

• sr_index (Tensor)

Return type:

Tensor

class mstar.model.voxcpm2.components.audio_vae.Snake1d(channels)

Bases: Module

Parameters:

channels (int)

forward(x)

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

mstar.model.voxcpm2.components.audio_vae.decoder_receptive_field_frames(decoder)

Latent frames of left context a frame’s audio needs to equal a whole decode.

Walks the decoder backwards from a frame’s first output sample; exact for the modules modelled in _earliest_input (20 frames for VoxCPM2). The sample-rate conditioning is per channel, so it adds nothing.

Parameters:

decoder (CausalDecoder)

Return type:

int

mstar.model.voxcpm2.components.audio_vae.fold_weight_norm(state)

weight_g / weight_v pairs -> weight (torch.nn.utils.weight_norm, dim 0).

Parameters:

state (dict[str, Tensor])

Return type:

dict[str, Tensor]