mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-05 10:00:28 +08:00
444 lines
16 KiB
Python
444 lines
16 KiB
Python
# MiniMax H3 audio VAE: DAC-lineage waveform encoder + BigVGAN decoder.
|
|
# Weight-norm parametrizations are folded into plain conv weights, so this
|
|
# module uses ordinary ops.Conv1d / ops.ConvTranspose1d and loads the converted
|
|
# checkpoint (plain "*.weight" tensors) with strict=True.
|
|
#
|
|
# Lineage / licenses of the reference implementation:
|
|
# DAC encoder: descript-audio-codec (MIT)
|
|
# BigVGAN decoder: NVIDIA BigVGAN (MIT), adapted from hifi-gan (MIT)
|
|
# Alias-free ops: junjun3518/alias-free-torch (Apache-2.0), julius (MIT)
|
|
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
import comfy.ops
|
|
|
|
ops = comfy.ops.disable_weight_init
|
|
|
|
|
|
# Snake activations
|
|
|
|
def snake(x, alpha, beta):
|
|
# x + 1/beta * sin^2(alpha * x)
|
|
t = torch.sin(alpha * x)
|
|
return t.mul_(t).mul_((beta + 1e-9).reciprocal()).add_(x)
|
|
|
|
|
|
class Snake1d(nn.Module):
|
|
"""Snake activation with per-channel alpha (encoder side)."""
|
|
|
|
def __init__(self, channels):
|
|
super().__init__()
|
|
self.alpha = nn.Parameter(torch.empty(1, channels, 1))
|
|
|
|
def forward(self, x):
|
|
alpha = comfy.ops.cast_to_input(self.alpha, x)
|
|
return snake(x, alpha, alpha)
|
|
|
|
|
|
class SnakeBeta(nn.Module):
|
|
"""SnakeBeta := x + 1/beta * sin^2(alpha * x); alpha/beta stored in log scale."""
|
|
|
|
def __init__(self, in_features):
|
|
super().__init__()
|
|
self.alpha = nn.Parameter(torch.empty(in_features))
|
|
self.beta = nn.Parameter(torch.empty(in_features))
|
|
|
|
def forward(self, x):
|
|
alpha = torch.exp(comfy.ops.cast_to_input(self.alpha, x)).view(1, -1, 1)
|
|
beta = torch.exp(comfy.ops.cast_to_input(self.beta, x)).view(1, -1, 1)
|
|
return snake(x, alpha, beta)
|
|
|
|
|
|
# Alias-free (anti-aliased) activation: kaiser-windowed sinc resampling
|
|
|
|
def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
|
|
# returns filter [1, 1, kernel_size]
|
|
even = kernel_size % 2 == 0
|
|
half_size = kernel_size // 2
|
|
|
|
# kaiser window design
|
|
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)
|
|
|
|
if even:
|
|
time = torch.arange(-half_size, half_size) + 0.5
|
|
else:
|
|
time = torch.arange(kernel_size) - half_size
|
|
|
|
filter_ = 2 * cutoff * window * torch.sinc(2 * cutoff * time)
|
|
# Normalize filter to have sum = 1, otherwise there is a small leakage of
|
|
# the constant component in the input signal.
|
|
filter_ /= filter_.sum()
|
|
return filter_.view(1, 1, kernel_size)
|
|
|
|
|
|
class UpSample1d(nn.Module):
|
|
def __init__(self, ratio=2, kernel_size=12):
|
|
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(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=kernel_size),
|
|
)
|
|
|
|
def forward(self, x):
|
|
_, C, _ = x.shape
|
|
x = F.pad(x, (self.pad, self.pad), mode="replicate")
|
|
x = F.conv_transpose1d(x, self.filter.expand(C, -1, -1).to(x.dtype), stride=self.stride, groups=C).mul_(self.ratio)
|
|
x = x[..., self.pad_left:-self.pad_right]
|
|
return x
|
|
|
|
|
|
class LowPassFilter1d(nn.Module):
|
|
def __init__(self, cutoff=0.5, half_width=0.6, stride=1, kernel_size=12):
|
|
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))
|
|
|
|
def forward(self, x):
|
|
_, C, _ = x.shape
|
|
x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate")
|
|
return F.conv1d(x, self.filter.expand(C, -1, -1).to(x.dtype), stride=self.stride, groups=C)
|
|
|
|
|
|
class DownSample1d(nn.Module):
|
|
def __init__(self, ratio=2, kernel_size=12):
|
|
super().__init__()
|
|
self.ratio = ratio
|
|
self.kernel_size = kernel_size
|
|
self.lowpass = LowPassFilter1d(
|
|
cutoff=0.5 / ratio,
|
|
half_width=0.6 / ratio,
|
|
stride=ratio,
|
|
kernel_size=self.kernel_size,
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.lowpass(x)
|
|
|
|
|
|
class Activation1d(nn.Module):
|
|
"""upsample x2 -> pointwise activation -> downsample x2 (anti-aliased)."""
|
|
|
|
def __init__(self, activation, up_ratio=2, down_ratio=2, up_kernel_size=12, down_kernel_size=12):
|
|
super().__init__()
|
|
self.act = activation
|
|
self.upsample = UpSample1d(up_ratio, up_kernel_size)
|
|
self.downsample = DownSample1d(down_ratio, down_kernel_size)
|
|
|
|
def forward(self, x):
|
|
x = self.upsample(x)
|
|
x = self.act(x)
|
|
x = self.downsample(x)
|
|
return x
|
|
|
|
|
|
# DAC encoder
|
|
|
|
class ResidualUnit(nn.Module):
|
|
def __init__(self, dim=16, dilation=1):
|
|
super().__init__()
|
|
pad = ((7 - 1) * dilation) // 2
|
|
self.block = nn.Sequential(
|
|
Snake1d(dim),
|
|
ops.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
|
|
Snake1d(dim),
|
|
ops.Conv1d(dim, dim, kernel_size=1),
|
|
)
|
|
|
|
def forward(self, x):
|
|
y = self.block(x)
|
|
pad = (x.shape[-1] - y.shape[-1]) // 2
|
|
if pad > 0:
|
|
x = x[..., pad:-pad]
|
|
return y.add_(x)
|
|
|
|
|
|
class EncoderBlock(nn.Module):
|
|
def __init__(self, dim=16, stride=1):
|
|
super().__init__()
|
|
self.block = nn.Sequential(
|
|
ResidualUnit(dim // 2, dilation=1),
|
|
ResidualUnit(dim // 2, dilation=3),
|
|
ResidualUnit(dim // 2, dilation=9),
|
|
Snake1d(dim // 2),
|
|
ops.Conv1d(
|
|
dim // 2,
|
|
dim,
|
|
kernel_size=2 * stride,
|
|
stride=stride,
|
|
padding=math.ceil(stride / 2),
|
|
),
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.block(x)
|
|
|
|
|
|
class Encoder(nn.Module):
|
|
def __init__(self, d_model=64, strides=(2, 4, 4, 5, 5), d_latent=2048):
|
|
super().__init__()
|
|
block = [ops.Conv1d(1, d_model, kernel_size=7, padding=3)]
|
|
for stride in strides:
|
|
d_model *= 2
|
|
block += [EncoderBlock(d_model, stride=stride)]
|
|
block += [
|
|
Snake1d(d_model),
|
|
ops.Conv1d(d_model, d_latent, kernel_size=3, padding=1),
|
|
]
|
|
self.block = nn.Sequential(*block)
|
|
|
|
def forward(self, x):
|
|
return self.block(x)
|
|
|
|
|
|
# Attention projection (encoder posterior head)
|
|
|
|
class GeGluMlp(nn.Module):
|
|
def __init__(self, in_features, hidden_features):
|
|
super().__init__()
|
|
self.norm = ops.LayerNorm(in_features)
|
|
self.act = nn.GELU(approximate="tanh")
|
|
self.w0 = ops.Linear(in_features, hidden_features)
|
|
self.w1 = ops.Linear(in_features, hidden_features)
|
|
self.w2 = ops.Linear(hidden_features, in_features)
|
|
|
|
def forward(self, x):
|
|
x = self.norm(x)
|
|
return self.w2(self.act(self.w0(x)).mul_(self.w1(x)))
|
|
|
|
|
|
class CausalAttention(nn.Module):
|
|
def __init__(self, in_dim, out_dim, num_heads):
|
|
super().__init__()
|
|
self.head_dim = in_dim // num_heads
|
|
self.num_heads = num_heads
|
|
self.out_dim = out_dim
|
|
self.qkv = ops.Linear(in_dim, in_dim * 3, bias=False)
|
|
self.q_bias = nn.Parameter(torch.empty(in_dim))
|
|
self.v_bias = nn.Parameter(torch.empty(in_dim))
|
|
self.register_buffer("zero_k_bias", torch.empty(in_dim))
|
|
self.proj = ops.Linear(out_dim, out_dim)
|
|
|
|
def forward(self, x):
|
|
B, N, C = x.shape
|
|
weight, _, offload_stream = comfy.ops.cast_bias_weight(self.qkv, x, offloadable=True)
|
|
qkv = F.linear(x, weight=weight, bias=comfy.ops.cast_to_input(torch.cat((self.q_bias, self.zero_k_bias, self.v_bias)), x))
|
|
comfy.ops.uncast_bias_weight(self.qkv, weight, None, offload_stream)
|
|
q, k, v = qkv.reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0)
|
|
|
|
# mean over heads then pool down to the latent width (in_dim >> out_dim)
|
|
x = comfy.ops.scaled_dot_product_attention(q, k, v, is_causal=True)
|
|
x = F.adaptive_avg_pool1d(torch.mean(x, dim=1), self.out_dim)
|
|
return self.proj(x)
|
|
|
|
|
|
class AttnProjection(nn.Module):
|
|
def __init__(self, in_dim, out_dim, num_heads, mlp_ratio=2):
|
|
super().__init__()
|
|
self.norm1 = ops.LayerNorm(in_dim)
|
|
self.attn = CausalAttention(in_dim, out_dim, num_heads)
|
|
self.proj = ops.Linear(in_dim, out_dim)
|
|
self.norm3 = ops.LayerNorm(in_dim)
|
|
|
|
self.norm2 = ops.LayerNorm(out_dim)
|
|
hidden_dim = int(out_dim * mlp_ratio)
|
|
self.mlp = GeGluMlp(in_features=out_dim, hidden_features=hidden_dim)
|
|
|
|
def forward(self, x):
|
|
# x: [B, T, in_dim]
|
|
x = self.proj(self.norm3(x)).add_(self.attn(self.norm1(x)))
|
|
return x.add_(self.mlp(self.norm2(x)))
|
|
|
|
|
|
# BigVGAN decoder
|
|
|
|
def get_padding(kernel_size, dilation=1):
|
|
return int((kernel_size * dilation - dilation) / 2)
|
|
|
|
|
|
class AMPBlock1(nn.Module):
|
|
def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
|
|
super().__init__()
|
|
self.convs1 = nn.ModuleList(
|
|
[
|
|
ops.Conv1d(channels, channels, kernel_size, stride=1, dilation=d, padding=get_padding(kernel_size, d))
|
|
for d in dilation
|
|
]
|
|
)
|
|
self.convs2 = nn.ModuleList(
|
|
[
|
|
ops.Conv1d(channels, channels, kernel_size, stride=1, dilation=1, padding=get_padding(kernel_size, 1))
|
|
for _ in range(len(dilation))
|
|
]
|
|
)
|
|
self.num_layers = len(self.convs1) + len(self.convs2)
|
|
self.activations = nn.ModuleList(
|
|
[Activation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)]
|
|
)
|
|
|
|
def forward(self, x):
|
|
acts1, acts2 = self.activations[::2], self.activations[1::2]
|
|
for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2):
|
|
xt = a1(x)
|
|
xt = c1(xt)
|
|
xt = a2(xt)
|
|
xt = c2(xt)
|
|
x = xt.add_(x)
|
|
return x
|
|
|
|
|
|
class BigVGAN(nn.Module):
|
|
"""BigVGAN vocoder (MiniMax H3 32 kHz configuration).
|
|
|
|
use_bias_at_final=False, use_tanh_at_final=False (output clamped to [-1, 1]).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
num_mels=2048,
|
|
upsample_initial_channel=1024,
|
|
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)),
|
|
):
|
|
super().__init__()
|
|
self.num_kernels = len(resblock_kernel_sizes)
|
|
self.num_upsamples = len(upsample_rates)
|
|
|
|
self.conv_pre = ops.Conv1d(num_mels, upsample_initial_channel, 7, 1, padding=3)
|
|
|
|
self.ups = nn.ModuleList()
|
|
for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
|
|
self.ups.append(
|
|
nn.ModuleList(
|
|
[
|
|
ops.ConvTranspose1d(
|
|
upsample_initial_channel // (2 ** i),
|
|
upsample_initial_channel // (2 ** (i + 1)),
|
|
k,
|
|
u,
|
|
padding=(k - u) // 2,
|
|
)
|
|
]
|
|
)
|
|
)
|
|
|
|
self.resblocks = nn.ModuleList()
|
|
for i in range(len(self.ups)):
|
|
ch = upsample_initial_channel // (2 ** (i + 1))
|
|
for k, d in zip(resblock_kernel_sizes, resblock_dilation_sizes):
|
|
self.resblocks.append(AMPBlock1(ch, k, d))
|
|
|
|
self.activation_post = Activation1d(activation=SnakeBeta(ch))
|
|
self.conv_post = ops.Conv1d(ch, 1, 7, 1, padding=3, bias=False)
|
|
|
|
def forward(self, x):
|
|
x = self.conv_pre(x)
|
|
|
|
for i in range(self.num_upsamples):
|
|
for i_up in range(len(self.ups[i])):
|
|
x = self.ups[i][i_up](x)
|
|
xs = None
|
|
for j in range(self.num_kernels):
|
|
if xs is None:
|
|
xs = self.resblocks[i * self.num_kernels + j](x)
|
|
else:
|
|
xs += self.resblocks[i * self.num_kernels + j](x)
|
|
x = xs.div_(self.num_kernels)
|
|
|
|
x = self.activation_post(x)
|
|
return self.conv_post(x).clamp_(-1.0, 1.0)
|
|
|
|
|
|
# Top-level VAE
|
|
|
|
class MiniMaxH3AudioVAE(nn.Module):
|
|
"""MiniMax H3 stereo audio VAE at 32 kHz.
|
|
|
|
Latents are [B, 32, 2, T]: 32 channels, 2 stereo channels, T frames at
|
|
40 latent frames per second (800 audio samples per latent frame). The
|
|
stereo channels are processed independently by the mono encoder/decoder.
|
|
Latents are normalized with the stored per-channel latents_mean/std.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
encoder_dim=64,
|
|
encoder_rates=(2, 4, 4, 5, 5),
|
|
latent_dim=2048,
|
|
decoder_dim=1024,
|
|
vae_latent_channels=32,
|
|
):
|
|
super().__init__()
|
|
self.sample_rate = 32000
|
|
|
|
self.hop_length = 1
|
|
for r in encoder_rates:
|
|
self.hop_length *= r
|
|
self.samples_per_latent = self.hop_length # 800
|
|
self.latents_per_second = self.sample_rate // self.hop_length # 40
|
|
self.output_sample_rate = self.sample_rate # read by LTXVAudioVAEDecode
|
|
|
|
self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim)
|
|
|
|
self.pre_block = AttnProjection(latent_dim, vae_latent_channels, num_heads=8)
|
|
|
|
self.mean_proj = ops.Conv1d(vae_latent_channels, vae_latent_channels, 1)
|
|
# logs_proj exists in the checkpoint but is unused at inference
|
|
# (encode returns the posterior mean, no sampling).
|
|
self.logs_proj = ops.Conv1d(vae_latent_channels, vae_latent_channels, 1)
|
|
|
|
self.dec_in_proj = ops.Conv1d(vae_latent_channels, latent_dim, 1)
|
|
self.decoder = BigVGAN(num_mels=latent_dim, upsample_initial_channel=decoder_dim)
|
|
|
|
self.register_buffer("latents_mean", torch.empty(vae_latent_channels))
|
|
self.register_buffer("latents_std", torch.empty(vae_latent_channels))
|
|
|
|
def decode(self, z):
|
|
"""Decode normalized latents [B, 32, 2, T] to stereo waveforms [B, 2, L] at 32 kHz."""
|
|
b, c, s, t = z.shape
|
|
z = z.permute(0, 2, 1, 3).reshape(b * s, c, t)
|
|
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
|
|
x = self.dec_in_proj(z)
|
|
x = self.decoder(x) # [b * s, 1, L], already clamped to [-1, 1]
|
|
return x.reshape(b, s, -1)
|
|
|
|
def encode(self, waveform):
|
|
"""Encode stereo waveforms [B, 2, L] at 32 kHz (in [-1, 1]) to normalized latents [B, 32, 2, T].
|
|
|
|
L is right-padded with zeros to a multiple of 800 samples; the returned
|
|
posterior mean is used directly (no sampling).
|
|
"""
|
|
b, s, length = waveform.shape
|
|
right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
|
|
waveform = F.pad(waveform, (0, right_pad))
|
|
x = waveform.reshape(b * s, 1, -1)
|
|
x = self.encoder(x) # [b * s, latent_dim, T]
|
|
x = self.pre_block(x.transpose(1, 2)).transpose(1, 2) # [b * s, 32, T]
|
|
z = self.mean_proj(x)
|
|
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 - mean) / std
|
|
return z.reshape(b, s, z.shape[1], z.shape[2]).permute(0, 2, 1, 3)
|