295 lines
11 KiB
Python
295 lines
11 KiB
Python
# 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]
|