192 lines
9.3 KiB
Python
192 lines
9.3 KiB
Python
|
"""Standalone decoder-only MiniMax H3 audio VAE."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import math
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from safetensors import safe_open
|
||
|
|
from torch import nn
|
||
|
|
from torch.nn import functional as F
|
||
|
|
|
||
|
|
|
||
|
|
def snake(x: torch.Tensor, alpha: torch.Tensor, beta: torch.Tensor) -> torch.Tensor:
|
||
|
|
t = torch.sin(alpha * x)
|
||
|
|
return t.mul_(t).mul_((beta + 1e-9).reciprocal()).add_(x)
|
||
|
|
|
||
|
|
|
||
|
|
class SnakeBeta(nn.Module):
|
||
|
|
def __init__(self, channels: int, *, device=None):
|
||
|
|
super().__init__()
|
||
|
|
self.alpha = nn.Parameter(torch.empty(channels, device=device))
|
||
|
|
self.beta = nn.Parameter(torch.empty(channels, device=device))
|
||
|
|
|
||
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
|
|
alpha = torch.exp(self.alpha.to(device=x.device, dtype=x.dtype)).view(1, -1, 1)
|
||
|
|
beta = torch.exp(self.beta.to(device=x.device, dtype=x.dtype)).view(1, -1, 1)
|
||
|
|
return snake(x, alpha, beta)
|
||
|
|
|
||
|
|
|
||
|
|
def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
|
||
|
|
even = kernel_size % 2 == 0
|
||
|
|
half_size = kernel_size // 2
|
||
|
|
delta_f = 4 * half_width
|
||
|
|
a = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
|
||
|
|
if a > 50.0:
|
||
|
|
beta = 0.1102 * (a - 8.7)
|
||
|
|
elif a >= 21.0:
|
||
|
|
beta = 0.5842 * (a - 21) ** 0.4 + 0.07886 * (a - 21.0)
|
||
|
|
else:
|
||
|
|
beta = 0.0
|
||
|
|
window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
|
||
|
|
time = torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size
|
||
|
|
filt = 2 * cutoff * window * torch.sinc(2 * cutoff * time)
|
||
|
|
filt /= filt.sum()
|
||
|
|
return filt.view(1, 1, kernel_size)
|
||
|
|
|
||
|
|
|
||
|
|
class UpSample1d(nn.Module):
|
||
|
|
def __init__(self, ratio: int = 2, kernel_size: int = 12, *, device=None):
|
||
|
|
super().__init__()
|
||
|
|
self.ratio = ratio
|
||
|
|
self.stride = ratio
|
||
|
|
self.pad = kernel_size // ratio - 1
|
||
|
|
self.pad_left = self.pad * ratio + (kernel_size - ratio) // 2
|
||
|
|
self.pad_right = self.pad * ratio + (kernel_size - ratio + 1) // 2
|
||
|
|
self.register_buffer("filter", kaiser_sinc_filter1d(0.5 / ratio, 0.6 / ratio, kernel_size).to(device=device))
|
||
|
|
|
||
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
|
|
_, channels, _ = x.shape
|
||
|
|
x = F.pad(x, (self.pad, self.pad), mode="replicate")
|
||
|
|
filt = self.filter.to(device=x.device, dtype=x.dtype).expand(channels, -1, -1)
|
||
|
|
x = F.conv_transpose1d(x, filt, stride=self.stride, groups=channels).mul_(self.ratio)
|
||
|
|
return x[..., self.pad_left:-self.pad_right]
|
||
|
|
|
||
|
|
|
||
|
|
class LowPassFilter1d(nn.Module):
|
||
|
|
def __init__(self, cutoff: float = 0.5, half_width: float = 0.6, stride: int = 1, kernel_size: int = 12, *, device=None):
|
||
|
|
super().__init__()
|
||
|
|
self.pad_left = kernel_size // 2 - int(kernel_size % 2 == 0)
|
||
|
|
self.pad_right = kernel_size // 2
|
||
|
|
self.stride = stride
|
||
|
|
self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size).to(device=device))
|
||
|
|
|
||
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
|
|
_, channels, _ = x.shape
|
||
|
|
x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate")
|
||
|
|
filt = self.filter.to(device=x.device, dtype=x.dtype).expand(channels, -1, -1)
|
||
|
|
return F.conv1d(x, filt, stride=self.stride, groups=channels)
|
||
|
|
|
||
|
|
|
||
|
|
class DownSample1d(nn.Module):
|
||
|
|
def __init__(self, ratio: int = 2, kernel_size: int = 12, *, device=None):
|
||
|
|
super().__init__()
|
||
|
|
self.lowpass = LowPassFilter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, stride=ratio, kernel_size=kernel_size, device=device)
|
||
|
|
|
||
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
|
|
return self.lowpass(x)
|
||
|
|
|
||
|
|
|
||
|
|
class Activation1d(nn.Module):
|
||
|
|
def __init__(self, activation: nn.Module, *, device=None):
|
||
|
|
super().__init__()
|
||
|
|
self.act = activation
|
||
|
|
self.upsample = UpSample1d(device=device)
|
||
|
|
self.downsample = DownSample1d(device=device)
|
||
|
|
|
||
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
|
|
return self.downsample(self.act(self.upsample(x)))
|
||
|
|
|
||
|
|
|
||
|
|
def get_padding(kernel_size: int, dilation: int = 1) -> int:
|
||
|
|
return int((kernel_size * dilation - dilation) / 2)
|
||
|
|
|
||
|
|
|
||
|
|
class AMPBlock1(nn.Module):
|
||
|
|
def __init__(self, channels: int, kernel_size: int = 3, dilation: tuple[int, int, int] = (1, 3, 5), *, device=None):
|
||
|
|
super().__init__()
|
||
|
|
self.convs1 = nn.ModuleList([nn.Conv1d(channels, channels, kernel_size, dilation=d, padding=get_padding(kernel_size, d), device=device) for d in dilation])
|
||
|
|
self.convs2 = nn.ModuleList([nn.Conv1d(channels, channels, kernel_size, dilation=1, padding=get_padding(kernel_size, 1), device=device) for _ in dilation])
|
||
|
|
self.activations = nn.ModuleList([Activation1d(SnakeBeta(channels, device=device), device=device) for _ in range(len(dilation) * 2)])
|
||
|
|
|
||
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
|
|
acts1, acts2 = self.activations[::2], self.activations[1::2]
|
||
|
|
for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, acts1, acts2):
|
||
|
|
residual = conv2(act2(conv1(act1(x))))
|
||
|
|
x = residual.add_(x)
|
||
|
|
return x
|
||
|
|
|
||
|
|
|
||
|
|
class BigVGAN(nn.Module):
|
||
|
|
def __init__(self, num_mels: int = 2048, upsample_initial_channel: int = 1024, *, device=None):
|
||
|
|
super().__init__()
|
||
|
|
upsample_rates = (5, 5, 2, 2, 2, 2, 2)
|
||
|
|
upsample_kernel_sizes = (9, 9, 4, 4, 4, 4, 4)
|
||
|
|
resblock_kernel_sizes = (3, 7, 11)
|
||
|
|
resblock_dilation_sizes = ((1, 3, 5), (1, 3, 5), (1, 3, 5))
|
||
|
|
self.num_kernels = len(resblock_kernel_sizes)
|
||
|
|
self.num_upsamples = len(upsample_rates)
|
||
|
|
self.conv_pre = nn.Conv1d(num_mels, upsample_initial_channel, 7, 1, padding=3, device=device)
|
||
|
|
self.ups = nn.ModuleList()
|
||
|
|
for index, (rate, kernel) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
|
||
|
|
self.ups.append(nn.ModuleList([nn.ConvTranspose1d(upsample_initial_channel // (2 ** index), upsample_initial_channel // (2 ** (index + 1)), kernel, rate, padding=(kernel - rate) // 2, device=device)]))
|
||
|
|
self.resblocks = nn.ModuleList()
|
||
|
|
for index in range(len(self.ups)):
|
||
|
|
channels = upsample_initial_channel // (2 ** (index + 1))
|
||
|
|
for kernel, dilation in zip(resblock_kernel_sizes, resblock_dilation_sizes):
|
||
|
|
self.resblocks.append(AMPBlock1(channels, kernel, dilation, device=device))
|
||
|
|
self.activation_post = Activation1d(SnakeBeta(channels, device=device), device=device)
|
||
|
|
self.conv_post = nn.Conv1d(channels, 1, 7, 1, padding=3, bias=False, device=device)
|
||
|
|
|
||
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
|
|
x = self.conv_pre(x)
|
||
|
|
for index in range(self.num_upsamples):
|
||
|
|
for upsample in self.ups[index]:
|
||
|
|
x = upsample(x)
|
||
|
|
combined = None
|
||
|
|
for kernel_index in range(self.num_kernels):
|
||
|
|
value = self.resblocks[index * self.num_kernels + kernel_index](x)
|
||
|
|
combined = value if combined is None else combined + value
|
||
|
|
x = combined.div_(self.num_kernels)
|
||
|
|
return self.conv_post(self.activation_post(x)).clamp_(-1.0, 1.0)
|
||
|
|
|
||
|
|
|
||
|
|
class MiniMaxH3AudioVAE(nn.Module):
|
||
|
|
"""Decoder-only MiniMax H3 stereo audio VAE at 32 kHz."""
|
||
|
|
|
||
|
|
def __init__(self, latent_channels: int = 32, latent_dim: int = 2048, decoder_dim: int = 1024, *, device=None):
|
||
|
|
super().__init__()
|
||
|
|
self.sample_rate = 32000
|
||
|
|
self.output_sample_rate = self.sample_rate
|
||
|
|
self.samples_per_latent = 800
|
||
|
|
self.latents_per_second = 40
|
||
|
|
self.dec_in_proj = nn.Conv1d(latent_channels, latent_dim, 1, device=device)
|
||
|
|
self.decoder = BigVGAN(num_mels=latent_dim, upsample_initial_channel=decoder_dim, device=device)
|
||
|
|
self.register_buffer("latents_mean", torch.empty(latent_channels, device=device))
|
||
|
|
self.register_buffer("latents_std", torch.empty(latent_channels, device=device))
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", dtype: torch.dtype = torch.float32) -> "MiniMaxH3AudioVAE":
|
||
|
|
model = cls(device="meta")
|
||
|
|
expected = model.state_dict()
|
||
|
|
with safe_open(str(path), framework="pt", device=str(device)) as checkpoint:
|
||
|
|
missing = sorted(set(expected) - set(checkpoint.keys()))
|
||
|
|
shape_errors = [(name, tuple(expected[name].shape), tuple(checkpoint.get_slice(name).get_shape())) for name in expected if name in checkpoint.keys() and tuple(expected[name].shape) != tuple(checkpoint.get_slice(name).get_shape())]
|
||
|
|
if missing or shape_errors:
|
||
|
|
details = ([f"missing: {', '.join(missing)}"] if missing else []) + ([f"shape mismatch: {shape_errors}"] if shape_errors else [])
|
||
|
|
raise ValueError("incompatible H3 audio VAE checkpoint; " + "; ".join(details))
|
||
|
|
weights = {name: checkpoint.get_tensor(name).to(dtype=dtype) for name in expected}
|
||
|
|
model.load_state_dict(weights, strict=True, assign=True)
|
||
|
|
return model
|
||
|
|
|
||
|
|
def decode(self, z: torch.Tensor) -> torch.Tensor:
|
||
|
|
"""Decode normalized latents `[B,32,2,T]` to waveform `[B,2,L]`."""
|
||
|
|
batch, channels, stereo, steps = z.shape
|
||
|
|
z = z.permute(0, 2, 1, 3).reshape(batch * stereo, channels, steps)
|
||
|
|
mean = self.latents_mean.view(1, -1, 1).to(device=z.device, dtype=z.dtype)
|
||
|
|
std = self.latents_std.view(1, -1, 1).to(device=z.device, dtype=z.dtype)
|
||
|
|
z = z * std + mean
|
||
|
|
waveform = self.decoder(self.dec_in_proj(z))
|
||
|
|
return waveform.reshape(batch, stereo, -1)
|