# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """DAC-style vocoder for MiniMax Music 3. Only the decoder half of the autoencoder ships in the checkpoint, which is all inference needs: the flow-matching stage produces the latent directly, so the encoder and the quantizer never run. Stereo is folded into the batch. A ``[B, 128, T]`` latent is read as two 64-channel mono latents per item, decoded as ``2B`` independent mono streams, and unfolded again at the end. The four transposed-convolution blocks upsample by 8, 8, 4 and 2, so one latent frame becomes 512 samples at 44.1 kHz. Module names match the ``vocoder/`` component folder key for key, so its state dict loads with ``strict=True`` and no remapping. The convolutions keep the checkpoint's weight-norm parameterization (``weight_g``/``weight_v``) through loading; :func:`remove_weight_norm` folds it away afterwards, since the factorization is fixed once the weights stop training. :class:`StreamingResampler` carries the 44.1 kHz output down to the 32 kHz response rate across chunk boundaries, which is not something a plain per-chunk resample can do. """ from __future__ import annotations import math from collections.abc import Sequence import torch from torch import Tensor, nn from vllm.logger import init_logger from .constants import DAV_SAMPLE_RATE, DIT_LATENT_CHANNELS, OUTPUT_CHANNELS, OUTPUT_SAMPLE_RATE logger = init_logger(__name__) __all__ = ["MiniMaxMusic3Vocoder", "StreamingResampler", "remove_weight_norm", "snake"] # Fixed by vocoder/config.json. _DECODER_INPUT_DIM = 1_024 _DECODER_HIDDEN_DIM = 1_536 _UPSAMPLING_RATIOS = (8, 8, 4, 2) _RESIDUAL_DILATIONS = (1, 3, 9) _KERNEL_SIZE = 7 # Each vocoder stream sees half the latent channels; the other half is the # opposite stereo channel. _STREAM_CHANNELS = DIT_LATENT_CHANNELS // OUTPUT_CHANNELS def snake(x: Tensor, alpha: Tensor) -> Tensor: """Periodic ``x + sin^2(alpha x) / alpha`` activation. The reshape to three dimensions is what the reference does and is a no-op for the rank-3 activations used here; it is kept so the numerics are identical for any rank. """ shape = x.shape flat = x.reshape(shape[0], shape[1], -1) flat = flat + (alpha + 1e-9).reciprocal() * torch.sin(alpha * flat).pow(2) return flat.reshape(shape) class Snake1d(nn.Module): """Snake activation with one learned frequency per channel.""" def __init__(self, channels: int) -> None: super().__init__() self.alpha = nn.Parameter(torch.ones(1, channels, 1)) def forward(self, x: Tensor) -> Tensor: return snake(x, self.alpha) def _wn_conv(*args: object, **kwargs: object) -> nn.Module: """A weight-normalized ``Conv1d``, named ``weight_g``/``weight_v``.""" return nn.utils.weight_norm(nn.Conv1d(*args, **kwargs)) # type: ignore[arg-type] def _wn_conv_transpose(*args: object, **kwargs: object) -> nn.Module: """A weight-normalized ``ConvTranspose1d``, named ``weight_g``/``weight_v``.""" return nn.utils.weight_norm(nn.ConvTranspose1d(*args, **kwargs)) # type: ignore[arg-type] def remove_weight_norm(module: nn.Module) -> int: """Fold weight normalization into every convolution weight in place. Returns: How many submodules were folded. """ removed = 0 for child in module.modules(): try: nn.utils.remove_weight_norm(child) except (ValueError, RuntimeError): continue removed += 1 return removed class ResidualUnit(nn.Module): """Dilated snake-gated residual block; preserves length and width.""" def __init__(self, dim: int, dilation: int) -> None: super().__init__() pad = (_KERNEL_SIZE - 1) * dilation // 2 self.snake1 = Snake1d(dim) self.conv1 = _wn_conv(dim, dim, kernel_size=_KERNEL_SIZE, dilation=dilation, padding=pad) self.snake2 = Snake1d(dim) self.conv2 = _wn_conv(dim, dim, kernel_size=1) def forward(self, x: Tensor) -> Tensor: y = self.conv2(self.snake2(self.conv1(self.snake1(x)))) if y.shape[-1] != x.shape[-1]: # Defensive centre crop. The padding above makes this a no-op for # every dilation the checkpoint ships. pad = (x.shape[-1] - y.shape[-1]) // 2 x = x[..., pad : x.shape[-1] - pad] return x + y class DecoderBlock(nn.Module): """One upsampling stage: activate, stretch by ``stride``, then refine.""" def __init__(self, input_dim: int, output_dim: int, stride: int) -> None: super().__init__() self.snake1 = Snake1d(input_dim) self.conv_t1 = _wn_conv_transpose( input_dim, output_dim, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2), ) self.res_unit1 = ResidualUnit(output_dim, _RESIDUAL_DILATIONS[0]) self.res_unit2 = ResidualUnit(output_dim, _RESIDUAL_DILATIONS[1]) self.res_unit3 = ResidualUnit(output_dim, _RESIDUAL_DILATIONS[2]) def forward(self, x: Tensor) -> Tensor: x = self.conv_t1(self.snake1(x)) return self.res_unit3(self.res_unit2(self.res_unit1(x))) class MiniMaxMusic3Vocoder(nn.Module): """Latent to waveform: ``[B, 128, T] -> [B, 2, T * 512]`` at 44.1 kHz.""" def __init__(self) -> None: super().__init__() self.dec_in_proj = nn.Conv1d(_STREAM_CHANNELS, _DECODER_INPUT_DIM, kernel_size=1) self.conv_in = _wn_conv( _DECODER_INPUT_DIM, _DECODER_HIDDEN_DIM, kernel_size=_KERNEL_SIZE, padding=_KERNEL_SIZE // 2, ) blocks: list[nn.Module] = [] output_dim = _DECODER_HIDDEN_DIM for index, stride in enumerate(_UPSAMPLING_RATIOS): input_dim = _DECODER_HIDDEN_DIM // (2**index) output_dim = _DECODER_HIDDEN_DIM // (2 ** (index + 1)) blocks.append(DecoderBlock(input_dim, output_dim, stride)) self.blocks = nn.ModuleList(blocks) self.snake_out = Snake1d(output_dim) self.conv_out = _wn_conv(output_dim, 1, kernel_size=_KERNEL_SIZE, padding=_KERNEL_SIZE // 2) @property def upsampling_factor(self) -> int: """Waveform samples produced per latent frame.""" return math.prod(_UPSAMPLING_RATIOS) @torch.inference_mode() def forward(self, latent: Tensor) -> Tensor: """Decode a latent to interleaved stereo. Args: latent: ``[B, 128, T]`` vocoder latent. Returns: ``[B, 2, T * 512]`` waveform at 44.1 kHz, unclamped. Raises: ValueError: If the latent is not ``[B, 128, T]``. """ if latent.ndim != 3 or latent.shape[1] != DIT_LATENT_CHANNELS: raise ValueError( f"MiniMax Music 3 vocoder latent must be [B,{DIT_LATENT_CHANNELS},T], got {tuple(latent.shape)}" ) bsz, _, frames = latent.shape folded = latent.reshape(bsz * OUTPUT_CHANNELS, _STREAM_CHANNELS, frames) h = self.conv_in(self.dec_in_proj(folded)) for block in self.blocks: h = block(h) wave = torch.tanh(self.conv_out(self.snake_out(h))) return wave.reshape(bsz, OUTPUT_CHANNELS, -1) class StreamingResampler: """Resample a chunked stream to the response rate without a seam. A polyphase resampler is time-invariant but not offset-invariant: it only reproduces the same output grid when the input offset is a whole multiple of ``orig // gcd(orig, new)`` samples, and it zero-pads both ends of every call. Resampling each decoded chunk on its own therefore shifts the grid by a fraction of a sample at every join and tapers the signal across it. Measured against a single whole-stream resample on 20 s of broadband noise, that costs a peak error of 2.2 (full scale) and drifts the length. This buffers instead. Only whole blocks are emitted, so every boundary lands on an exact output sample; the previously emitted tail is prepended as real left context and the not-yet-emitted head appended as real right context, and both are dropped from the result. On the same measurement the emitted stream is bit-identical to the whole-stream resample. One instance per request. It holds at most a couple of blocks of audio. """ def __init__( self, *, orig_freq: int = DAV_SAMPLE_RATE, new_freq: int = OUTPUT_SAMPLE_RATE, channels: int = OUTPUT_CHANNELS, context_blocks: int = 1, margin_blocks: int = 1, ) -> None: self.orig_freq = int(orig_freq) self.new_freq = int(new_freq) self.channels = int(channels) # The output grid repeats every this many input samples. self._block = self.orig_freq // math.gcd(self.orig_freq, self.new_freq) self._context_len = self._block * context_blocks self._margin = self._block * margin_blocks self._pending: Tensor | None = None self._context: Tensor | None = None def _empty(self, like: Tensor | None) -> Tensor: if like is None: return torch.zeros((self.channels, 0), dtype=torch.float32) return like.new_zeros((self.channels, 0)) def reset(self) -> None: """Drop all buffered audio.""" self._pending = None self._context = None @torch.inference_mode() def push(self, chunks: Sequence[Tensor], *, flush: bool = False) -> Tensor: """Append source-rate audio and return whatever is now emittable. Args: chunks: New ``[channels, T]`` segments at ``orig_freq``, in order. flush: Emit the remainder too, because the stream has ended. Returns: ``[channels, T]`` at ``new_freq``; empty when nothing is ready. """ import torchaudio.functional as audio_functional parts = [chunk for chunk in chunks if chunk.numel()] if self._pending is not None: parts.insert(0, self._pending) self._pending = None if not parts: return self._empty(self._context) buffered = torch.cat(parts, dim=-1) if len(parts) > 1 else parts[0] total = int(buffered.shape[-1]) if flush: emit_len, lookahead = total, 0 else: emit_len = max(0, ((total - self._margin) // self._block) * self._block) lookahead = min(self._margin, total - emit_len) if emit_len == 0: self._pending = buffered return self._empty(buffered) padded = buffered[..., : emit_len + lookahead] start = 0 if self._context is not None: padded = torch.cat((self._context, padded), dim=-1) # Exact because the context is a whole number of blocks. start = int(self._context.shape[-1]) * self.new_freq // self.orig_freq resampled = audio_functional.resample(padded.float(), self.orig_freq, self.new_freq) # Round up, matching how a whole-stream resample sizes its output. This # only bites on the flush, where the remainder is not a whole block; # every other emission divides exactly. stop = start + -(-emit_len * self.new_freq // self.orig_freq) if flush: self.reset() else: self._context = buffered[..., max(0, emit_len - self._context_len) : emit_len].clone() self._pending = buffered[..., emit_len:].clone() if emit_len < total else None return resampled[..., start:stop]