861 lines
30 KiB
Python
861 lines
30 KiB
Python
# Copyright 2026 Tencent.
|
|
# token2wav: audio token codes -> waveform (inference only)
|
|
# Pipeline: Token -> Latent (flow matching) -> Waveform (BigVGAN)
|
|
|
|
import math
|
|
from collections import OrderedDict, namedtuple
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from einops import rearrange
|
|
from torch.nn.utils import weight_norm
|
|
|
|
from vllm_omni.model_executor.models.common.alias_free_activation import AliasFreeActivation1d
|
|
from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta
|
|
|
|
# --------------- Utilities ---------------
|
|
|
|
|
|
class JsonHParams:
|
|
def __init__(self, **kwargs):
|
|
for k, v in kwargs.items():
|
|
if isinstance(v, dict):
|
|
v = JsonHParams(**v)
|
|
if isinstance(v, str) and v.lower() in ["non", "none", "nil", "null"]:
|
|
v = None
|
|
self[k] = v
|
|
|
|
def keys(self):
|
|
return self.__dict__.keys()
|
|
|
|
def items(self):
|
|
return self.__dict__.items()
|
|
|
|
def values(self):
|
|
return self.__dict__.values()
|
|
|
|
def __len__(self):
|
|
return len(self.__dict__)
|
|
|
|
def __getitem__(self, key):
|
|
return getattr(self, key)
|
|
|
|
def __setitem__(self, key, value):
|
|
return setattr(self, key, value)
|
|
|
|
def __contains__(self, key):
|
|
return key in self.__dict__
|
|
|
|
def __repr__(self):
|
|
return self.__dict__.__repr__()
|
|
|
|
def get(self, key, default=None):
|
|
return self.__dict__.get(key, default)
|
|
|
|
|
|
def exists(val):
|
|
return val is not None
|
|
|
|
|
|
def eval_decorator(fn):
|
|
def inner(model, *args, **kwargs):
|
|
was_training = model.training
|
|
model.eval()
|
|
out = fn(model, *args, **kwargs)
|
|
model.train(was_training)
|
|
return out
|
|
|
|
return inner
|
|
|
|
|
|
def init_weights(m, mean=0.0, std=0.01):
|
|
classname = m.__class__.__name__
|
|
if classname.find("Conv") != -1:
|
|
m.weight.data.normal_(mean, std)
|
|
|
|
|
|
def get_padding(kernel_size, dilation=1):
|
|
return int((kernel_size * dilation - dilation) / 2)
|
|
|
|
|
|
def modulate(x, shift, scale):
|
|
return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
|
|
|
|
|
|
def rotate_half(x):
|
|
x1, x2 = x.chunk(2, dim=-1)
|
|
return torch.cat((-x2, x1), dim=-1)
|
|
|
|
|
|
@torch.autocast(enabled=False, device_type="cuda")
|
|
def apply_rotary_pos_emb(pos, t):
|
|
return t * pos.cos() + rotate_half(t) * pos.sin()
|
|
|
|
|
|
# sinc, kaiser_sinc_filter1d removed — now in common.alias_free_activation
|
|
|
|
|
|
# --------------- Basic Layers ---------------
|
|
|
|
|
|
class Dropout(nn.Module):
|
|
def __init__(self, p: float = 0.5, inplace: bool = False, force_drop: bool = False, **kwargs):
|
|
super().__init__()
|
|
if p < 0.0 or p > 1.0:
|
|
raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}")
|
|
self.p = p
|
|
self.inplace = inplace
|
|
self.force_drop = force_drop
|
|
|
|
def forward(self, x, **kwargs):
|
|
return F.dropout(x, p=self.p, training=True if self.force_drop else self.training, inplace=self.inplace)
|
|
|
|
|
|
class EmbeddingTable(nn.Embedding):
|
|
def __init__(self, num_embeddings: int, embedding_dim: int, pad_id=-1, **kwargs):
|
|
super().__init__(num_embeddings, embedding_dim, **kwargs)
|
|
nn.init.normal_(self.weight, 0.0, embedding_dim**-0.5)
|
|
self.pad_id = pad_id
|
|
self.output_dim = embedding_dim
|
|
|
|
def forward(self, x):
|
|
if self.pad_id is not None:
|
|
mask = x == self.pad_id
|
|
x = x.masked_fill(mask, 0)
|
|
outputs = super().forward(x)
|
|
if self.pad_id is not None:
|
|
outputs = outputs.masked_fill(mask.unsqueeze(-1), 0.0)
|
|
return outputs
|
|
|
|
|
|
class Linear(nn.Linear):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
bias: bool = True,
|
|
w_init_gain: str = "linear",
|
|
activation=None,
|
|
**kwargs,
|
|
):
|
|
super().__init__(in_channels, out_channels, bias=bias)
|
|
self.activation = activation if activation is not None else nn.Identity()
|
|
self.output_dim = out_channels
|
|
if w_init_gain is not None:
|
|
if isinstance(w_init_gain, str):
|
|
gain = nn.init.calculate_gain(w_init_gain)
|
|
else:
|
|
gain = w_init_gain
|
|
nn.init.xavier_uniform_(self.weight, gain=gain)
|
|
if bias:
|
|
nn.init.constant_(self.bias, 0.0)
|
|
|
|
def forward(self, x, **kwargs):
|
|
return self.activation(super().forward(x))
|
|
|
|
|
|
class Conv1d(nn.Conv1d):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
kernel_size: int = 1,
|
|
stride: int = 1,
|
|
dilation: int = 1,
|
|
groups: int = 1,
|
|
padding_mode: str = "zeros",
|
|
bias: bool = True,
|
|
padding=None,
|
|
causal: bool = False,
|
|
bn: bool = False,
|
|
activation=None,
|
|
w_init_gain=None,
|
|
input_transpose: bool = False,
|
|
**kwargs,
|
|
):
|
|
self.causal = causal
|
|
if padding is None:
|
|
if causal:
|
|
padding = 0
|
|
self.left_padding = dilation * (kernel_size - 1)
|
|
else:
|
|
padding = get_padding(kernel_size, dilation)
|
|
|
|
super().__init__(
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size,
|
|
stride=stride,
|
|
padding=padding,
|
|
dilation=dilation,
|
|
groups=groups,
|
|
padding_mode=padding_mode,
|
|
bias=bias,
|
|
)
|
|
|
|
self.in_channels = in_channels
|
|
self.transpose = input_transpose
|
|
self.bn = nn.BatchNorm1d(out_channels) if bn else nn.Identity()
|
|
self.activation = activation if activation is not None else nn.Identity()
|
|
if w_init_gain is not None:
|
|
nn.init.xavier_uniform_(self.weight, gain=nn.init.calculate_gain(w_init_gain))
|
|
|
|
def forward(self, x):
|
|
if self.transpose or x.size(1) != self.in_channels:
|
|
assert x.size(2) == self.in_channels
|
|
x = x.transpose(1, 2)
|
|
self.transpose = True
|
|
|
|
if self.causal:
|
|
x = F.pad(x.unsqueeze(2), (self.left_padding, 0, 0, 0)).squeeze(2)
|
|
|
|
outputs = self.activation(self.bn(super().forward(x)))
|
|
return outputs.transpose(1, 2) if self.transpose else outputs
|
|
|
|
|
|
class ConvTranspose1d(nn.ConvTranspose1d):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
kernel_size: int,
|
|
stride: int = 1,
|
|
output_padding: int = 0,
|
|
groups: int = 1,
|
|
bias: bool = True,
|
|
dilation: int = 1,
|
|
padding=None,
|
|
padding_mode: str = "zeros",
|
|
causal: bool = False,
|
|
input_transpose: bool = False,
|
|
**kwargs,
|
|
):
|
|
if padding is None:
|
|
padding = 0 if causal else (kernel_size - stride) // 2
|
|
if causal:
|
|
assert padding == 0, "padding is not allowed in causal ConvTranspose1d."
|
|
assert kernel_size == 2 * stride, "kernel_size must be equal to 2*stride in Causal ConvTranspose1d."
|
|
|
|
super().__init__(
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size,
|
|
stride=stride,
|
|
padding=padding,
|
|
output_padding=output_padding,
|
|
groups=groups,
|
|
bias=bias,
|
|
dilation=dilation,
|
|
padding_mode=padding_mode,
|
|
)
|
|
|
|
self.causal = causal
|
|
self.stride = stride
|
|
self.transpose = input_transpose
|
|
|
|
def forward(self, x):
|
|
if self.transpose or x.size(1) != self.in_channels:
|
|
assert x.size(2) == self.in_channels
|
|
x = x.transpose(1, 2)
|
|
self.transpose = True
|
|
|
|
x = super().forward(x)
|
|
if self.causal:
|
|
x = x[:, :, : -self.stride]
|
|
return x.transpose(1, 2) if self.transpose else x
|
|
|
|
|
|
class ConvPositionEmbed(nn.Module):
|
|
def __init__(self, hidden_size: int, kernel_size: int = 31, groups: int = 16):
|
|
super().__init__()
|
|
self.conv = Conv1d(
|
|
hidden_size, hidden_size, kernel_size, groups=groups, input_transpose=True, activation=nn.GELU()
|
|
)
|
|
|
|
def forward(self, x, mask=None):
|
|
if mask is not None:
|
|
mask = mask[..., None]
|
|
x = x.masked_fill(~mask, 0.0)
|
|
x = self.conv(x)
|
|
if mask is not None:
|
|
x = x.masked_fill(~mask, 0.0)
|
|
return x
|
|
|
|
|
|
# --------------- Anti-aliasing (BigVGAN) ---------------
|
|
|
|
|
|
# LowPassFilter1d, UpSample1d, DownSample1d, Activation1d removed —
|
|
# now provided by common.alias_free_activation.AliasFreeActivation1d
|
|
|
|
|
|
# --------------- Activations ---------------
|
|
|
|
|
|
# SnakeBeta removed — now in common.snake_activation (with alpha_logscale support)
|
|
|
|
|
|
# --------------- Attention ---------------
|
|
|
|
AttentionConfig = namedtuple("AttentionConfig", ["enable_flash", "enable_math", "enable_mem_efficient"])
|
|
|
|
|
|
class RotaryEmbedding(nn.Module):
|
|
def __init__(self, dim, theta=50000):
|
|
super().__init__()
|
|
inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
|
|
self.register_buffer("inv_freq", inv_freq)
|
|
|
|
@property
|
|
def device(self):
|
|
return self.inv_freq.device
|
|
|
|
@torch.autocast(enabled=False, device_type="cuda")
|
|
def forward(self, t):
|
|
if not torch.is_tensor(t):
|
|
t = torch.arange(t, device=self.device)
|
|
t = t.type_as(self.inv_freq)
|
|
freqs = torch.einsum("i , j -> i j", t, self.inv_freq)
|
|
freqs = torch.cat((freqs, freqs), dim=-1)
|
|
return freqs
|
|
|
|
|
|
class MultiHeadAttention(nn.Module):
|
|
def __init__(
|
|
self,
|
|
hidden_size: int,
|
|
num_heads: int = 8,
|
|
qkv_bias: bool = False,
|
|
qk_norm: bool = False,
|
|
attn_drop: float = 0.0,
|
|
dropout: float = 0.0,
|
|
max_position_embeddings: int = 4096,
|
|
norm_layer: nn.Module = nn.LayerNorm,
|
|
rotary_bias: bool = False,
|
|
**kwargs,
|
|
):
|
|
super().__init__()
|
|
assert hidden_size % num_heads == 0, "hidden_size should be divisible by num_heads"
|
|
self.num_heads = num_heads
|
|
self.head_dim = hidden_size // num_heads
|
|
self.scale = self.head_dim**-0.5
|
|
self.max_position_embeddings = max_position_embeddings
|
|
self.rotary_bias = rotary_bias
|
|
|
|
self.q_proj = Linear(hidden_size, hidden_size, bias=qkv_bias)
|
|
self.k_proj = Linear(hidden_size, hidden_size, bias=qkv_bias)
|
|
self.v_proj = Linear(hidden_size, hidden_size, bias=qkv_bias)
|
|
self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
|
|
self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
|
|
|
|
self.attn_drop = Dropout(attn_drop)
|
|
self.o_proj = Linear(hidden_size, hidden_size)
|
|
self.o_dropout = Dropout(dropout)
|
|
|
|
self.cpu_config = AttentionConfig(True, True, True)
|
|
device_properties = torch.cuda.get_device_properties(torch.device("cuda"))
|
|
if device_properties.major == 8 and device_properties.minor == 0:
|
|
self.cuda_config = AttentionConfig(True, True, True)
|
|
else:
|
|
self.cuda_config = AttentionConfig(False, True, True)
|
|
|
|
if self.rotary_bias:
|
|
self.rotary = RotaryEmbedding(self.head_dim)
|
|
|
|
def forward(self, q, k=None, v=None, mask=None):
|
|
k = k or q
|
|
v = v or q
|
|
B, L, C = q.shape
|
|
B, S, C = v.shape
|
|
if mask is not None:
|
|
if mask.ndim == 2:
|
|
assert L == S
|
|
mask = rearrange(mask, "b j -> b 1 1 j")
|
|
mask = mask.expand(-1, self.num_heads, L, -1)
|
|
elif mask.ndim == 3:
|
|
assert mask.size(1) == L and mask.size(2) == S
|
|
mask = mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1)
|
|
|
|
q, k, v = self.q_proj(q), self.k_proj(k), self.v_proj(v)
|
|
q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads)
|
|
k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads)
|
|
v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads)
|
|
q, k = self.q_norm(q), self.k_norm(k)
|
|
|
|
config = self.cuda_config if q.is_cuda else self.cpu_config
|
|
attn_bias = torch.zeros(B, self.num_heads, L, S, dtype=q.dtype, device=q.device)
|
|
|
|
if self.rotary_bias:
|
|
if L == S:
|
|
rotary_emb = self.rotary(L)
|
|
q, k = map(lambda x: apply_rotary_pos_emb(rotary_emb, x), (q, k))
|
|
else:
|
|
q_rotary_emb = self.rotary(L)
|
|
k_rotary_emb = self.rotary(S)
|
|
q = apply_rotary_pos_emb(q_rotary_emb, q)
|
|
k = apply_rotary_pos_emb(k_rotary_emb, k)
|
|
|
|
if mask is not None:
|
|
attn_bias.masked_fill_(mask.logical_not(), float("-inf"))
|
|
|
|
with torch.backends.cuda.sdp_kernel(**config._asdict()):
|
|
out = F.scaled_dot_product_attention(
|
|
q, k, v, attn_mask=attn_bias, dropout_p=self.attn_drop.p if self.training else 0.0
|
|
)
|
|
|
|
out = rearrange(out, "b h n d -> b n (h d)")
|
|
out = self.o_dropout(self.o_proj(out))
|
|
return out
|
|
|
|
|
|
# --------------- DiT Modules ---------------
|
|
|
|
|
|
class TimestepEmbedder(nn.Module):
|
|
def __init__(self, hidden_size, frequency_embedding_size=256):
|
|
super().__init__()
|
|
self.mlp = nn.Sequential(
|
|
Linear(frequency_embedding_size, hidden_size, bias=True),
|
|
nn.SiLU(),
|
|
Linear(hidden_size, hidden_size, bias=True),
|
|
)
|
|
self.frequency_embedding_size = frequency_embedding_size
|
|
|
|
@staticmethod
|
|
def timestep_embedding(t, dim, max_period=10000):
|
|
half = dim // 2
|
|
freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
|
|
device=t.device
|
|
)
|
|
args = t[:, None].float() * freqs[None]
|
|
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
|
if dim % 2:
|
|
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
|
return embedding
|
|
|
|
def forward(self, t):
|
|
t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
|
|
t_emb = self.mlp(t_freq)
|
|
return t_emb
|
|
|
|
|
|
class FinalLayer(nn.Module):
|
|
def __init__(self, hidden_size, output_size):
|
|
super().__init__()
|
|
self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
|
|
self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
|
|
self.linear = Linear(hidden_size, output_size, bias=True)
|
|
|
|
def forward(self, x, c, mask=None):
|
|
shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
|
|
x = modulate(self.norm(x), shift, scale)
|
|
x = self.linear(x)
|
|
return x
|
|
|
|
|
|
# --------------- Transformer ---------------
|
|
|
|
|
|
class Mlp(nn.Module):
|
|
def __init__(self, hidden_size, ffn_hidden_size=4096, act_layer=nn.GELU, dropout=0.0, **kwargs):
|
|
super().__init__()
|
|
self.fc1 = Linear(hidden_size, ffn_hidden_size)
|
|
self.act = act_layer()
|
|
self.fc2 = Linear(ffn_hidden_size, hidden_size)
|
|
self.drop = Dropout(dropout)
|
|
|
|
def forward(self, x, **kwargs):
|
|
x = self.fc1(x)
|
|
x = self.act(x)
|
|
x = self.drop(x)
|
|
x = self.fc2(x)
|
|
x = self.drop(x)
|
|
return x
|
|
|
|
|
|
class TransformerBlock(nn.Module):
|
|
def __init__(
|
|
self,
|
|
attention: nn.Module,
|
|
ffn: nn.Module,
|
|
hidden_size: int = 1024,
|
|
modulation: bool = False,
|
|
eps: float = 1e-6,
|
|
**kwargs,
|
|
):
|
|
super().__init__()
|
|
self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=not modulation, eps=eps)
|
|
self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=not modulation, eps=eps)
|
|
self.attn = attention
|
|
self.ffn = ffn
|
|
self.modulation = modulation
|
|
if modulation:
|
|
self.modulation_layer = nn.Sequential(nn.SiLU(), Linear(hidden_size, 6 * hidden_size, bias=True))
|
|
nn.init.constant_(self.modulation_layer[-1].weight, 0.0)
|
|
nn.init.constant_(self.modulation_layer[-1].bias, 0.0)
|
|
|
|
def forward(self, x, condition=None, mask=None):
|
|
if condition is None:
|
|
assert not self.modulation, "Without global condition, must set modulation to False"
|
|
else:
|
|
assert self.modulation, "With global condition, must set modulation to True"
|
|
shift_attn, scale_attn, gate_attn, shift_ffn, scale_ffn, gate_ffn = self.modulation_layer(condition).chunk(
|
|
6, dim=1
|
|
)
|
|
|
|
if condition is not None:
|
|
x = x + gate_attn.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_attn, scale_attn), mask=mask)
|
|
else:
|
|
x = x + self.attn(self.norm1(x), mask=mask)
|
|
|
|
if condition is not None:
|
|
x = x + gate_ffn.unsqueeze(1) * self.ffn(modulate(self.norm2(x), shift_ffn, scale_ffn))
|
|
else:
|
|
x = x + self.ffn(self.norm2(x), mask=mask)
|
|
return x
|
|
|
|
|
|
# --------------- Flow Matching: Token -> Latent ---------------
|
|
|
|
|
|
class Token2latentFlowMatching(nn.Module):
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
self.config = config
|
|
self.model_dim = config.transformer.hidden_size
|
|
self.token_input_dim = config.get("token_input_dim", self.model_dim)
|
|
self.target_dim = config.z_dim
|
|
|
|
self.spkr_embed_dim = config.get("spkr_embed_dim", 512)
|
|
self.cond_proj = Linear(self.model_dim + self.spkr_embed_dim, self.model_dim)
|
|
self.token_pad_id = -1
|
|
if config.upsample_factor > 1:
|
|
self.token_proj = nn.Sequential(
|
|
Linear(self.token_input_dim, self.model_dim, bias=True),
|
|
ConvTranspose1d(
|
|
self.model_dim,
|
|
self.model_dim,
|
|
stride=config.upsample_factor,
|
|
kernel_size=config.upsample_factor * 2,
|
|
),
|
|
)
|
|
else:
|
|
self.token_proj = Linear(self.token_input_dim, self.model_dim)
|
|
|
|
self.transformer_input_proj = Linear(self.model_dim + self.target_dim * 2, self.model_dim)
|
|
self.time_embedder = TimestepEmbedder(self.model_dim)
|
|
self.conv_embed = ConvPositionEmbed(hidden_size=self.model_dim, kernel_size=31, groups=16)
|
|
self.blocks = nn.ModuleList()
|
|
for _ in range(config.transformer.num_layers):
|
|
attn_block = MultiHeadAttention(**config.transformer)
|
|
ffn_block = Mlp(act_layer=lambda: nn.GELU(approximate="tanh"), **config.transformer)
|
|
self.blocks.append(TransformerBlock(attn_block, ffn_block, **config.transformer))
|
|
self.output_layer = FinalLayer(config.transformer.hidden_size, self.target_dim)
|
|
|
|
def cond_mask_spkr_embed(self, x, spkr_embed):
|
|
b, device = x.size(0), x.device
|
|
if not exists(spkr_embed):
|
|
spkr_embed = torch.zeros(b, self.spkr_embed_dim, device=device, dtype=x.dtype)
|
|
return spkr_embed
|
|
|
|
def vectorfield_forward(self, inputs, times, self_attn_mask, g_cond=None):
|
|
t = self.time_embedder(times)
|
|
if g_cond is None:
|
|
c = t
|
|
else:
|
|
cond_inp = torch.cat([t, g_cond], dim=-1)
|
|
c = self.cond_proj(cond_inp)
|
|
|
|
ut = self.transformer_input_proj(inputs)
|
|
ut = self.conv_embed(ut, mask=self_attn_mask) + ut
|
|
|
|
for block in self.blocks:
|
|
ut = block(ut, c, mask=self_attn_mask)
|
|
ut = self.output_layer(ut, c)
|
|
return ut
|
|
|
|
@eval_decorator
|
|
@torch.no_grad()
|
|
def inference(
|
|
self,
|
|
*,
|
|
token: torch.Tensor,
|
|
prefix_target: torch.Tensor | None = None,
|
|
spkr_embed: torch.Tensor | None = None,
|
|
s_steps: int | None = 10,
|
|
cfg_alpha: float | None = 2.0,
|
|
rescale_logits: bool = False,
|
|
**kwargs,
|
|
):
|
|
b, device = token.size(0), token.device
|
|
assert b == 1, "Only support batch_size == 1 when inference"
|
|
|
|
token = self.token_proj(token)
|
|
ge = self.cond_mask_spkr_embed(token, spkr_embed)
|
|
tgt_lens = token.size(1)
|
|
latent_cond = torch.zeros(b, tgt_lens, self.target_dim).to(device)
|
|
if prefix_target is not None:
|
|
latent_cond[:, : prefix_target.size(1), :] = prefix_target
|
|
sample, _ = self.sample(
|
|
tokens=token, audio=latent_cond, steps=s_steps, alpha=cfg_alpha, g_cond=ge, rescale_logits=rescale_logits
|
|
)
|
|
return sample
|
|
|
|
def sample(self, tokens, audio, steps, alpha=None, g_cond=None, rescale_logits=False):
|
|
noise = torch.randn([audio.size(0), audio.shape[1], self.target_dim], device=audio.device)
|
|
times = torch.linspace(0, 1, steps, device=audio.device)
|
|
|
|
def solver(t, z):
|
|
if alpha is None:
|
|
output = torch.cat([tokens, audio, z], dim=-1)
|
|
return self.vectorfield_forward(inputs=output, times=t.unsqueeze(0), self_attn_mask=None, g_cond=g_cond)
|
|
tokens_empty = torch.zeros(*audio.shape[:2], self.model_dim, device=tokens.device, dtype=tokens.dtype)
|
|
tokens_t = torch.cat([tokens_empty, tokens], dim=0)
|
|
audio_t = torch.cat([audio, audio], dim=0)
|
|
audio_noizy_t = torch.cat([z, z], dim=0)
|
|
t_t = torch.stack([t, t], dim=0)
|
|
c = g_cond
|
|
if g_cond is not None:
|
|
c = torch.cat([g_cond, g_cond], dim=0)
|
|
output = torch.cat([tokens_t, audio_t, audio_noizy_t], dim=-1)
|
|
predicted_mix = self.vectorfield_forward(inputs=output, times=t_t, self_attn_mask=None, g_cond=c)
|
|
predicted_conditioned = predicted_mix[1].unsqueeze(0)
|
|
predicted_unconditioned = predicted_mix[0].unsqueeze(0)
|
|
|
|
prediction = predicted_conditioned + (predicted_conditioned - predicted_unconditioned) * alpha
|
|
prediction_rescaled = predicted_conditioned.std() * (prediction / prediction.std())
|
|
if rescale_logits:
|
|
return prediction_rescaled
|
|
else:
|
|
return prediction
|
|
|
|
try:
|
|
from torchdiffeq import odeint
|
|
except ImportError:
|
|
raise ImportError("Covo-Audio code2wav requires `torchdiffeq`. Install it with: pip install torchdiffeq")
|
|
trajectory = odeint(solver, noise, times, atol=1e-5, rtol=1e-5, method="midpoint")
|
|
return trajectory[-1], trajectory
|
|
|
|
|
|
class Token2latentFlowMatchingWithEmbed(Token2latentFlowMatching):
|
|
def __init__(self, config, **kwargs):
|
|
super().__init__(config)
|
|
self.vocab_size = config.token_vocab_size
|
|
self.token_embedding = EmbeddingTable(
|
|
num_embeddings=self.vocab_size, embedding_dim=self.token_input_dim, pad_id=self.token_pad_id
|
|
)
|
|
|
|
def inference(
|
|
self, *, token, prefix_target=None, spkr_embed=None, s_steps=10, cfg_alpha=2, rescale_logits=False, **kwargs
|
|
):
|
|
token = self.token_embedding(token)
|
|
return super().inference(
|
|
token=token,
|
|
prefix_target=prefix_target,
|
|
spkr_embed=spkr_embed,
|
|
s_steps=s_steps,
|
|
cfg_alpha=cfg_alpha,
|
|
rescale_logits=rescale_logits,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
# --------------- Vocoder: Latent -> Waveform (BigVGAN) ---------------
|
|
|
|
|
|
class AMPBlock1(torch.nn.Module):
|
|
def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5), causal=True):
|
|
super().__init__()
|
|
self.h = h
|
|
|
|
self.convs1 = nn.ModuleList(
|
|
[
|
|
weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], causal=causal)),
|
|
weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], causal=causal)),
|
|
weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2], causal=causal)),
|
|
]
|
|
)
|
|
self.convs1.apply(init_weights)
|
|
|
|
self.convs2 = nn.ModuleList(
|
|
[
|
|
weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)),
|
|
weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)),
|
|
weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1, causal=causal)),
|
|
]
|
|
)
|
|
self.convs2.apply(init_weights)
|
|
|
|
self.num_layers = len(self.convs1) + len(self.convs2)
|
|
|
|
self.activations = nn.ModuleList(
|
|
[
|
|
AliasFreeActivation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale))
|
|
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 + x
|
|
return x
|
|
|
|
|
|
class BigVGANFlowVAE(nn.Module):
|
|
def __init__(self, h):
|
|
super().__init__()
|
|
self.config = h
|
|
self.h = h
|
|
causal = h.causal
|
|
self.hop_size = np.prod(h.downsample_rates)
|
|
|
|
self.num_kernels = len(h.resblock_kernel_sizes)
|
|
self.num_upsamples = len(h.upsample_rates)
|
|
|
|
self.conv_pre = weight_norm(Conv1d(h.latent_dim, h.upsample_initial_channel, 7, 1, causal=False))
|
|
|
|
self.ups = nn.ModuleList()
|
|
for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
|
|
self.ups.append(
|
|
nn.ModuleList(
|
|
[
|
|
weight_norm(
|
|
ConvTranspose1d(
|
|
h.upsample_initial_channel // (2**i),
|
|
h.upsample_initial_channel // (2 ** (i + 1)),
|
|
k,
|
|
u,
|
|
causal=causal,
|
|
)
|
|
)
|
|
]
|
|
)
|
|
)
|
|
|
|
self.resblocks = nn.ModuleList()
|
|
for i in range(len(self.ups)):
|
|
ch = h.upsample_initial_channel // (2 ** (i + 1))
|
|
for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
|
|
self.resblocks.append(AMPBlock1(h, ch, k, d, causal=causal))
|
|
|
|
activation_post = SnakeBeta(ch, alpha_logscale=h.snake_logscale)
|
|
self.activation_post = AliasFreeActivation1d(activation=activation_post)
|
|
|
|
self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, causal=causal))
|
|
|
|
for i in range(len(self.ups)):
|
|
self.ups[i].apply(init_weights)
|
|
self.conv_post.apply(init_weights)
|
|
|
|
def inference_from_latents(self, x, do_sample=True, noise_scale=1.0):
|
|
if self.h.use_vae and do_sample:
|
|
assert x.size(1) == self.h.latent_dim * 2, f"Input must be like [B, D, H], got {x.shape}"
|
|
m_q, logs_q = torch.split(x, self.h.latent_dim, dim=1)
|
|
x = m_q + torch.randn_like(m_q) * torch.exp(logs_q) * noise_scale
|
|
else:
|
|
assert x.size(1) == self.h.latent_dim, f"Input must be like [B, D, H], got {x.shape}"
|
|
|
|
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 / self.num_kernels
|
|
|
|
x = self.activation_post(x)
|
|
x = self.conv_post(x)
|
|
x = torch.tanh(x)
|
|
return x
|
|
|
|
|
|
# --------------- Token2WavDecoder ---------------
|
|
|
|
|
|
class Token2WavDecoder(nn.Module):
|
|
def __init__(self, config, **kwargs):
|
|
super().__init__()
|
|
self.config = config
|
|
|
|
self.wavegan = BigVGANFlowVAE(config.wavegan)
|
|
self.wavegan_hop_size = np.prod(config.wavegan.downsample_rates)
|
|
self.global_mean_var = getattr(config, "global_mean_var", None)
|
|
if self.global_mean_var is not None:
|
|
mean_var_data = torch.from_numpy(np.load(self.global_mean_var)).float().squeeze()
|
|
global_mean, global_var = mean_var_data.chunk(2, 0)
|
|
self.register_buffer("global_mean", global_mean)
|
|
self.register_buffer("global_var", global_var)
|
|
|
|
self.token2latent = Token2latentFlowMatchingWithEmbed(config.token2latent)
|
|
|
|
self.upsample_factor = self.token2latent.config.get("upsample_factor", 1)
|
|
self.wav_input_sr = config.get("wav_input_sr", 24000)
|
|
|
|
self.trainable_module = ["wavegan", "token2latent"]
|
|
|
|
def state_dict(self):
|
|
param_dict = OrderedDict()
|
|
for name in self.trainable_module:
|
|
state = self.get_submodule(name).state_dict(prefix=f"{name}.")
|
|
param_dict.update(state)
|
|
return param_dict
|
|
|
|
def load_state_dict(self, param_dict):
|
|
for name in self.trainable_module:
|
|
module_state = OrderedDict()
|
|
name_len = len(name)
|
|
for k, v in param_dict.items():
|
|
if k.startswith(f"{name}."):
|
|
new_k = k[name_len + 1 :]
|
|
module_state[new_k] = v
|
|
self.get_submodule(name).load_state_dict(module_state, strict=False)
|
|
|
|
@torch.no_grad()
|
|
def preprocess_infer_data(self, data):
|
|
zero_spkr = data.get("zero_spkr", False)
|
|
|
|
token = data["target_token"]
|
|
prefix_target = None
|
|
res_prefix_len = 0
|
|
prompt_token = data.get("prompt_token", None)
|
|
prompt_latent = data.get("prompt_latent", None)
|
|
if exists(prompt_token):
|
|
token = torch.cat([prompt_token, token], dim=1)
|
|
if exists(prompt_latent):
|
|
prefix_target = prompt_latent
|
|
res_prefix_len = prompt_latent.shape[1]
|
|
|
|
spkr_embed = None
|
|
if not zero_spkr:
|
|
spkr_embed = data.get("spkr_embed", None)
|
|
res = {"token": token, "prefix_target": prefix_target, "spkr_embed": spkr_embed}
|
|
return res, res_prefix_len
|
|
|
|
@eval_decorator
|
|
@torch.no_grad()
|
|
def inference(self, data, **kwargs):
|
|
infer_data, prefix_len = self.preprocess_infer_data(data)
|
|
prefix_target = infer_data["prefix_target"]
|
|
res_latents = self.token2latent.inference(**infer_data, **kwargs)
|
|
if self.global_mean_var is not None:
|
|
res_latents = self.global_mean + res_latents * torch.sqrt(self.global_var)
|
|
if exists(prefix_target):
|
|
prefix_target = self.global_mean + prefix_target * torch.sqrt(self.global_var)
|
|
if exists(prefix_target):
|
|
res_latents[:, : prefix_target.shape[1]] = prefix_target
|
|
res_latents = res_latents[:, prefix_len:]
|
|
audio = self.wavegan.inference_from_latents(res_latents.transpose(1, 2), do_sample=False)
|
|
return audio
|