NancyFyongandNancyFyong
e74fe3c6e8
Signed-off-by: NancyFyong <NancyFyong@users.noreply.github.com> Co-authored-by: NancyFyong <NancyFyong@users.noreply.github.com>
601 lines
23 KiB
Python
601 lines
23 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Audio8 TTS neural audio codec (44.1 kHz, 10 codebooks).
|
|
|
|
Vendored from the reference ``modeling_arktts_codec.py`` shipped with
|
|
``Audio8/Audio8-TTS-Preview-0.6b`` (Apache-2.0). The module tree and attribute
|
|
names are reproduced verbatim so ``codec.pth`` loads under ``strict=True``.
|
|
|
|
Inference-only deviations: no ``torch.jit.script`` on Snake (it fights
|
|
``torch.compile`` / CUDA graphs), RoPE cached in a plain dict instead of a
|
|
non-persistent buffer (the pattern that leaves the reference LM with
|
|
uninitialised RoPE under ``from_pretrained`` on transformers >= 5), and
|
|
:func:`build_arktts_codec` prunes the unused encoder or decoder half per stage.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch import Tensor
|
|
from torch.nn.utils import weight_norm as legacy_weight_norm
|
|
from torch.nn.utils.parametrizations import weight_norm
|
|
|
|
# The codec's own transformer blocks always use base 10000, independent of the
|
|
# 1e6 rope_base of the DualAR language model.
|
|
_CODEC_ROPE_BASE = 10000.0
|
|
|
|
|
|
def _rope_table(length: int, head_dim: int, base: float, device: torch.device) -> Tensor:
|
|
frequencies = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
|
|
phases = torch.outer(torch.arange(length, device=device), frequencies)
|
|
values = torch.polar(torch.ones_like(phases), phases)
|
|
return torch.stack((values.real, values.imag), dim=-1).to(torch.bfloat16)
|
|
|
|
|
|
def _apply_rope(x: Tensor, values: Tensor) -> Tensor:
|
|
shaped = x.float().reshape(*x.shape[:-1], -1, 2)
|
|
values = values.view(1, shaped.shape[1], 1, shaped.shape[3], 2)
|
|
output = torch.stack(
|
|
(
|
|
shaped[..., 0] * values[..., 0] - shaped[..., 1] * values[..., 1],
|
|
shaped[..., 1] * values[..., 0] + shaped[..., 0] * values[..., 1],
|
|
),
|
|
dim=-1,
|
|
)
|
|
return output.flatten(3).to(x.dtype)
|
|
|
|
|
|
class ArkttsCodecRMSNorm(nn.Module):
|
|
def __init__(self, dim: int, eps: float = 1e-5) -> None:
|
|
super().__init__()
|
|
self.eps = float(eps)
|
|
self.weight = nn.Parameter(torch.ones(dim))
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
output = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
|
|
return output.to(x.dtype) * self.weight
|
|
|
|
|
|
class ArkttsCodecLayerScale(nn.Module):
|
|
def __init__(self, dim: int, init_values: float = 1e-2) -> None:
|
|
super().__init__()
|
|
self.gamma = nn.Parameter(init_values * torch.ones(dim))
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
return x * self.gamma
|
|
|
|
|
|
class ArkttsCodecTransformerConfig:
|
|
"""Plain option holder for the codec's transformer blocks.
|
|
|
|
Deliberately not a dataclass / msgspec Struct: it mirrors a reference
|
|
dataclass whose only job is to carry constructor arguments, and it is
|
|
instantiated from Python literals inside this module only.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
n_layer: int,
|
|
n_head: int,
|
|
dim: int,
|
|
intermediate_size: int,
|
|
n_local_heads: int = -1,
|
|
head_dim: int = 64,
|
|
rope_base: float = _CODEC_ROPE_BASE,
|
|
norm_eps: float = 1e-5,
|
|
channels_first: bool = True,
|
|
) -> None:
|
|
self.n_layer = int(n_layer)
|
|
self.n_head = int(n_head)
|
|
self.dim = int(dim)
|
|
self.intermediate_size = int(intermediate_size)
|
|
self.n_local_heads = int(n_head) if n_local_heads == -1 else int(n_local_heads)
|
|
self.head_dim = int(head_dim)
|
|
self.rope_base = float(rope_base)
|
|
self.norm_eps = float(norm_eps)
|
|
self.channels_first = bool(channels_first)
|
|
|
|
|
|
class ArkttsCodecAttention(nn.Module):
|
|
def __init__(self, config: ArkttsCodecTransformerConfig) -> None:
|
|
super().__init__()
|
|
total = (config.n_head + 2 * config.n_local_heads) * config.head_dim
|
|
self.wqkv = nn.Linear(config.dim, total, bias=False)
|
|
self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False)
|
|
self.n_head = config.n_head
|
|
self.n_local_heads = config.n_local_heads
|
|
self.head_dim = config.head_dim
|
|
|
|
def forward(self, x: Tensor, rope_values: Tensor, mask: Tensor) -> Tensor:
|
|
batch, length, _ = x.shape
|
|
query_size = self.n_head * self.head_dim
|
|
kv_size = self.n_local_heads * self.head_dim
|
|
query, key, value = self.wqkv(x).split((query_size, kv_size, kv_size), dim=-1)
|
|
query = query.view(batch, length, self.n_head, self.head_dim)
|
|
key = key.view(batch, length, self.n_local_heads, self.head_dim)
|
|
value = value.view(batch, length, self.n_local_heads, self.head_dim)
|
|
query = _apply_rope(query, rope_values).transpose(1, 2)
|
|
key = _apply_rope(key, rope_values).transpose(1, 2)
|
|
value = value.transpose(1, 2)
|
|
output = F.scaled_dot_product_attention(
|
|
query,
|
|
key,
|
|
value,
|
|
attn_mask=mask,
|
|
enable_gqa=self.n_local_heads != self.n_head,
|
|
)
|
|
output = output.transpose(1, 2).contiguous().view(batch, length, query_size)
|
|
return self.wo(output)
|
|
|
|
|
|
class ArkttsCodecFeedForward(nn.Module):
|
|
def __init__(self, config: ArkttsCodecTransformerConfig) -> None:
|
|
super().__init__()
|
|
self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False)
|
|
self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False)
|
|
self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False)
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
return self.w2(F.silu(self.w1(x)) * self.w3(x))
|
|
|
|
|
|
class ArkttsCodecTransformerBlock(nn.Module):
|
|
def __init__(self, config: ArkttsCodecTransformerConfig) -> None:
|
|
super().__init__()
|
|
self.attention = ArkttsCodecAttention(config)
|
|
self.feed_forward = ArkttsCodecFeedForward(config)
|
|
self.ffn_norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps)
|
|
self.attention_norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps)
|
|
self.attention_layer_scale = ArkttsCodecLayerScale(config.dim)
|
|
self.ffn_layer_scale = ArkttsCodecLayerScale(config.dim)
|
|
|
|
def forward(self, x: Tensor, rope_values: Tensor, mask: Tensor) -> Tensor:
|
|
hidden = x + self.attention_layer_scale(self.attention(self.attention_norm(x), rope_values, mask))
|
|
return hidden + self.ffn_layer_scale(self.feed_forward(self.ffn_norm(hidden)))
|
|
|
|
|
|
class ArkttsCodecWindowTransformer(nn.Module):
|
|
"""Causal transformer with a sliding attention window."""
|
|
|
|
def __init__(
|
|
self,
|
|
config: ArkttsCodecTransformerConfig,
|
|
input_dim: int,
|
|
window_size: int | None,
|
|
) -> None:
|
|
super().__init__()
|
|
self.layers = nn.ModuleList([ArkttsCodecTransformerBlock(config) for _ in range(config.n_layer)])
|
|
self.norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps)
|
|
self.window_size = window_size
|
|
self.channels_first = config.channels_first
|
|
self.input_proj = nn.Linear(input_dim, config.dim) if input_dim != config.dim else nn.Identity()
|
|
self.output_proj = nn.Linear(config.dim, input_dim) if input_dim != config.dim else nn.Identity()
|
|
self.look_ahead_conv = nn.Identity()
|
|
self.head_dim = config.head_dim
|
|
self.rope_base = config.rope_base
|
|
self._mask_cache: dict[tuple[int, str], Tensor] = {}
|
|
self._rope_cache: dict[tuple[int, str], Tensor] = {}
|
|
|
|
def _mask(self, length: int, device: torch.device) -> Tensor:
|
|
key = (length, str(device))
|
|
cached = self._mask_cache.get(key)
|
|
if cached is not None:
|
|
return cached
|
|
row = torch.arange(length, device=device)[:, None]
|
|
column = torch.arange(length, device=device)[None, :]
|
|
mask = column <= row
|
|
if self.window_size is not None:
|
|
mask &= column >= (row - self.window_size + 1).clamp_min(0)
|
|
mask = mask[None, None]
|
|
self._mask_cache[key] = mask
|
|
return mask
|
|
|
|
def _rope(self, length: int, device: torch.device) -> Tensor:
|
|
key = (length, str(device))
|
|
cached = self._rope_cache.get(key)
|
|
if cached is not None:
|
|
return cached
|
|
values = _rope_table(length, self.head_dim, self.rope_base, device)
|
|
self._rope_cache[key] = values
|
|
return values
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
if self.channels_first:
|
|
x = x.transpose(1, 2)
|
|
x = self.look_ahead_conv(self.input_proj(x))
|
|
length = int(x.shape[1])
|
|
mask = self._mask(length, x.device)
|
|
rope_values = self._rope(length, x.device)
|
|
for layer in self.layers:
|
|
x = layer(x, rope_values, mask)
|
|
x = self.output_proj(self.norm(x))
|
|
return x.transpose(1, 2) if self.channels_first else x
|
|
|
|
|
|
def _extra_padding(x: Tensor, kernel_size: int, stride: int, padding_total: int = 0) -> int:
|
|
length = x.shape[-1]
|
|
frames = (length - kernel_size + padding_total) / stride + 1
|
|
ideal = (math.ceil(frames) - 1) * stride + kernel_size - padding_total
|
|
return ideal - length
|
|
|
|
|
|
class ArkttsCausalConv1d(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
kernel_size: int,
|
|
dilation: int = 1,
|
|
stride: int = 1,
|
|
groups: int = 1,
|
|
) -> None:
|
|
super().__init__()
|
|
self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, groups=groups)
|
|
self.stride = stride
|
|
self.kernel_size = (kernel_size - 1) * dilation + 1
|
|
self.padding = self.kernel_size - self.stride
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
right = _extra_padding(x, self.kernel_size, self.stride, self.padding)
|
|
return self.conv(F.pad(x, (self.padding, right))).contiguous()
|
|
|
|
def apply_weight_norm(self) -> ArkttsCausalConv1d:
|
|
self.conv = weight_norm(self.conv)
|
|
return self
|
|
|
|
|
|
class ArkttsCausalConvTranspose1d(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
kernel_size: int,
|
|
stride: int = 1,
|
|
dilation: int = 1,
|
|
) -> None:
|
|
super().__init__()
|
|
self.conv = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation)
|
|
self.stride = stride
|
|
self.kernel_size = kernel_size
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
x = self.conv(x)
|
|
crop = self.kernel_size - self.stride
|
|
return x[..., : x.shape[-1] - crop].contiguous() if crop else x.contiguous()
|
|
|
|
def apply_weight_norm(self) -> ArkttsCausalConvTranspose1d:
|
|
self.conv = weight_norm(self.conv)
|
|
return self
|
|
|
|
|
|
def _causal_wn_conv(*args: Any, **kwargs: Any) -> ArkttsCausalConv1d:
|
|
return ArkttsCausalConv1d(*args, **kwargs).apply_weight_norm()
|
|
|
|
|
|
def _causal_wn_transpose(*args: Any, **kwargs: Any) -> ArkttsCausalConvTranspose1d:
|
|
return ArkttsCausalConvTranspose1d(*args, **kwargs).apply_weight_norm()
|
|
|
|
|
|
class ArkttsSnake1d(nn.Module):
|
|
def __init__(self, channels: int) -> None:
|
|
super().__init__()
|
|
self.alpha = nn.Parameter(torch.ones(1, channels, 1))
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
shape = x.shape
|
|
x = x.reshape(shape[0], shape[1], -1)
|
|
x = x + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * x).pow(2)
|
|
return x.reshape(shape)
|
|
|
|
|
|
class ArkttsResidualUnit(nn.Module):
|
|
def __init__(self, dim: int, dilation: int) -> None:
|
|
super().__init__()
|
|
self.block = nn.Sequential(
|
|
ArkttsSnake1d(dim),
|
|
_causal_wn_conv(dim, dim, kernel_size=7, dilation=dilation),
|
|
ArkttsSnake1d(dim),
|
|
_causal_wn_conv(dim, dim, kernel_size=1),
|
|
)
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
output = self.block(x)
|
|
difference = x.shape[-1] - output.shape[-1]
|
|
if difference > 0:
|
|
x = x[..., :-difference]
|
|
return x + output
|
|
|
|
|
|
class ArkttsEncoderBlock(nn.Module):
|
|
def __init__(self, dim: int, stride: int, transformer_layers: int) -> None:
|
|
super().__init__()
|
|
modules: list[nn.Module] = [
|
|
ArkttsResidualUnit(dim // 2, 1),
|
|
ArkttsResidualUnit(dim // 2, 3),
|
|
ArkttsResidualUnit(dim // 2, 9),
|
|
ArkttsSnake1d(dim // 2),
|
|
_causal_wn_conv(dim // 2, dim, kernel_size=2 * stride, stride=stride),
|
|
]
|
|
if transformer_layers:
|
|
config = ArkttsCodecTransformerConfig(
|
|
n_layer=transformer_layers,
|
|
n_head=dim // 64,
|
|
dim=dim,
|
|
intermediate_size=dim * 3,
|
|
)
|
|
modules.append(ArkttsCodecWindowTransformer(config, dim, window_size=512))
|
|
else:
|
|
modules.append(nn.Identity())
|
|
self.block = nn.Sequential(*modules)
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
return self.block(x)
|
|
|
|
|
|
class ArkttsEncoder(nn.Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
dim = 64
|
|
modules: list[nn.Module] = [_causal_wn_conv(1, dim, kernel_size=7)]
|
|
for stride, transformer_layers in zip((2, 4, 8, 8), (0, 0, 0, 4), strict=True):
|
|
dim *= 2
|
|
modules.append(ArkttsEncoderBlock(dim, stride, transformer_layers))
|
|
modules.extend((ArkttsSnake1d(dim), _causal_wn_conv(dim, 1024, kernel_size=3)))
|
|
self.block = nn.Sequential(*modules)
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
return self.block(x)
|
|
|
|
|
|
class ArkttsDecoderBlock(nn.Module):
|
|
def __init__(self, input_dim: int, output_dim: int, stride: int) -> None:
|
|
super().__init__()
|
|
self.block = nn.Sequential(
|
|
ArkttsSnake1d(input_dim),
|
|
_causal_wn_transpose(input_dim, output_dim, kernel_size=2 * stride, stride=stride),
|
|
ArkttsResidualUnit(output_dim, 1),
|
|
ArkttsResidualUnit(output_dim, 3),
|
|
ArkttsResidualUnit(output_dim, 9),
|
|
)
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
return self.block(x)
|
|
|
|
|
|
class ArkttsDecoder(nn.Module):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
channels = 1536
|
|
modules: list[nn.Module] = [_causal_wn_conv(1024, channels, kernel_size=7)]
|
|
output_dim = channels
|
|
for index, stride in enumerate((8, 8, 4, 2)):
|
|
input_dim = channels // (2**index)
|
|
output_dim = channels // (2 ** (index + 1))
|
|
modules.append(ArkttsDecoderBlock(input_dim, output_dim, stride))
|
|
modules.extend(
|
|
(
|
|
ArkttsSnake1d(output_dim),
|
|
_causal_wn_conv(output_dim, 1, kernel_size=7),
|
|
nn.Tanh(),
|
|
)
|
|
)
|
|
self.model = nn.Sequential(*modules)
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
return self.model(x)
|
|
|
|
|
|
class ArkttsVectorQuantizer(nn.Module):
|
|
def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int) -> None:
|
|
super().__init__()
|
|
self.codebook_size = int(codebook_size)
|
|
self.codebook_dim = int(codebook_dim)
|
|
self.in_proj = legacy_weight_norm(nn.Conv1d(input_dim, codebook_dim, kernel_size=1))
|
|
self.out_proj = legacy_weight_norm(nn.Conv1d(codebook_dim, input_dim, kernel_size=1))
|
|
self.codebook = nn.Embedding(codebook_size, codebook_dim)
|
|
|
|
def decode_code(self, indices: Tensor) -> Tensor:
|
|
return F.embedding(indices, self.codebook.weight).transpose(1, 2)
|
|
|
|
def decode_latents(self, latents: Tensor) -> tuple[Tensor, Tensor]:
|
|
batch, _, length = latents.shape
|
|
flattened = F.normalize(latents.transpose(1, 2).reshape(batch * length, -1))
|
|
codebook = F.normalize(self.codebook.weight)
|
|
distances = (
|
|
flattened.pow(2).sum(1, keepdim=True)
|
|
- 2 * flattened @ codebook.t()
|
|
+ codebook.pow(2).sum(1, keepdim=True).t()
|
|
)
|
|
indices = (-distances).argmax(1).view(batch, length)
|
|
return self.decode_code(indices), indices
|
|
|
|
def forward(self, z: Tensor) -> tuple[Tensor, Tensor]:
|
|
projected = self.in_proj(z)
|
|
quantized, indices = self.decode_latents(projected)
|
|
return self.out_proj(quantized), indices
|
|
|
|
|
|
class ArkttsResidualQuantizer(nn.Module):
|
|
def __init__(self, input_dim: int, n_codebooks: int, codebook_size: int, codebook_dim: int) -> None:
|
|
super().__init__()
|
|
self.n_codebooks = int(n_codebooks)
|
|
self.codebook_size = int(codebook_size)
|
|
self.quantizers = nn.ModuleList(
|
|
[ArkttsVectorQuantizer(input_dim, codebook_size, codebook_dim) for _ in range(n_codebooks)]
|
|
)
|
|
|
|
def forward(self, z: Tensor) -> tuple[Tensor, Tensor]:
|
|
quantized_sum: Tensor | float = 0.0
|
|
residual = z
|
|
codes: list[Tensor] = []
|
|
for quantizer in self.quantizers:
|
|
quantized, indices = quantizer(residual)
|
|
quantized_sum = quantized_sum + quantized
|
|
residual = residual - quantized
|
|
codes.append(indices)
|
|
return quantized_sum, torch.stack(codes, dim=1)
|
|
|
|
def from_codes(self, codes: Tensor) -> Tensor:
|
|
output: Tensor | float = 0.0
|
|
for index in range(codes.shape[1]):
|
|
projected = self.quantizers[index].decode_code(codes[:, index])
|
|
output = output + self.quantizers[index].out_proj(projected)
|
|
return output
|
|
|
|
|
|
class ArkttsConvNeXtBlock(nn.Module):
|
|
def __init__(self, dim: int) -> None:
|
|
super().__init__()
|
|
self.dwconv = ArkttsCausalConv1d(dim, dim, kernel_size=7, groups=dim)
|
|
self.norm = nn.LayerNorm(dim, eps=1e-6)
|
|
self.pwconv1 = nn.Linear(dim, 4 * dim)
|
|
self.act = nn.GELU()
|
|
self.pwconv2 = nn.Linear(4 * dim, dim)
|
|
self.gamma = nn.Parameter(1e-6 * torch.ones(dim))
|
|
|
|
def forward(self, x: Tensor) -> Tensor:
|
|
residual = x
|
|
x = self.dwconv(x).permute(0, 2, 1)
|
|
x = self.pwconv2(self.act(self.pwconv1(self.norm(x))))
|
|
x = (self.gamma * x).permute(0, 2, 1)
|
|
return residual + x
|
|
|
|
|
|
class ArkttsDownsampleQuantizer(nn.Module):
|
|
"""1 semantic codebook (4096) + 9 residual codebooks (1024), 4x downsampled."""
|
|
|
|
def __init__(
|
|
self,
|
|
post_n_layer: int = 8,
|
|
post_n_head: int = 16,
|
|
post_n_local_heads: int = 8,
|
|
post_intermediate_size: int = 1216,
|
|
) -> None:
|
|
super().__init__()
|
|
self.semantic_quantizer = ArkttsResidualQuantizer(1024, 1, 4096, 8)
|
|
self.quantizer = ArkttsResidualQuantizer(1024, 9, 1024, 8)
|
|
self.downsample = nn.Sequential(
|
|
nn.Sequential(ArkttsCausalConv1d(1024, 1024, kernel_size=2, stride=2), ArkttsConvNeXtBlock(1024)),
|
|
nn.Sequential(ArkttsCausalConv1d(1024, 1024, kernel_size=2, stride=2), ArkttsConvNeXtBlock(1024)),
|
|
)
|
|
self.upsample = nn.Sequential(
|
|
nn.Sequential(ArkttsCausalConvTranspose1d(1024, 1024, kernel_size=2, stride=2), ArkttsConvNeXtBlock(1024)),
|
|
nn.Sequential(ArkttsCausalConvTranspose1d(1024, 1024, kernel_size=2, stride=2), ArkttsConvNeXtBlock(1024)),
|
|
)
|
|
self.pre_module = ArkttsCodecWindowTransformer(
|
|
ArkttsCodecTransformerConfig(n_layer=8, n_head=16, dim=1024, intermediate_size=3072),
|
|
1024,
|
|
window_size=128,
|
|
)
|
|
self.post_module = ArkttsCodecWindowTransformer(
|
|
ArkttsCodecTransformerConfig(
|
|
n_layer=post_n_layer,
|
|
n_head=post_n_head,
|
|
n_local_heads=post_n_local_heads,
|
|
dim=1024,
|
|
intermediate_size=post_intermediate_size,
|
|
),
|
|
1024,
|
|
window_size=128,
|
|
)
|
|
self.semantic_predictor_module = nn.Identity()
|
|
|
|
def forward(self, z: Tensor) -> tuple[Tensor, Tensor]:
|
|
original_length = z.shape[-1]
|
|
z = self.pre_module(self.downsample(z))
|
|
semantic, semantic_codes = self.semantic_quantizer(z)
|
|
residual, residual_codes = self.quantizer(z - semantic)
|
|
z = self.upsample(self.post_module(semantic + residual))
|
|
difference = original_length - z.shape[-1]
|
|
if difference > 0:
|
|
z = F.pad(z, (difference, 0))
|
|
elif difference < 0:
|
|
z = z[..., -difference:]
|
|
return z, torch.cat((semantic_codes, residual_codes), dim=1)
|
|
|
|
def decode(self, indices: Tensor) -> Tensor:
|
|
# Cloned because the clamps below are in-place and this is the caller's
|
|
# codes buffer; once per decode chunk, not per AR step.
|
|
indices = indices.clone()
|
|
indices[:, 0].clamp_(0, self.semantic_quantizer.codebook_size - 1)
|
|
indices[:, 1:].clamp_(0, self.quantizer.codebook_size - 1)
|
|
semantic = self.semantic_quantizer.from_codes(indices[:, :1])
|
|
residual = self.quantizer.from_codes(indices[:, 1:])
|
|
return self.upsample(self.post_module(semantic + residual))
|
|
|
|
|
|
class ArkttsCodec(nn.Module):
|
|
"""Audio8 TTS codec: waveform <-> 10 codebooks at ~21.5 frames/s."""
|
|
|
|
sample_rate = 44100
|
|
frame_length = 2048
|
|
|
|
def __init__(
|
|
self,
|
|
post_n_layer: int = 8,
|
|
post_n_head: int = 16,
|
|
post_n_local_heads: int = 8,
|
|
post_intermediate_size: int = 1216,
|
|
) -> None:
|
|
super().__init__()
|
|
self.encoder = ArkttsEncoder()
|
|
self.quantizer = ArkttsDownsampleQuantizer(
|
|
post_n_layer=post_n_layer,
|
|
post_n_head=post_n_head,
|
|
post_n_local_heads=post_n_local_heads,
|
|
post_intermediate_size=post_intermediate_size,
|
|
)
|
|
self.decoder = ArkttsDecoder()
|
|
|
|
@torch.inference_mode()
|
|
def encode(self, audio: Tensor, audio_lengths: Tensor | None = None) -> tuple[Tensor, Tensor]:
|
|
"""Encode ``[B, 1, samples]`` (or ``[B, samples]``) to ``[B, 10, frames]``."""
|
|
if audio.ndim == 2:
|
|
audio = audio[:, None]
|
|
if audio.ndim != 3 or audio.shape[1] != 1:
|
|
raise ValueError(f"audio must have shape [B, 1, samples], got {tuple(audio.shape)}")
|
|
original_length = int(audio.shape[-1])
|
|
right = math.ceil(original_length / self.frame_length) * self.frame_length - original_length
|
|
if right:
|
|
audio = F.pad(audio, (0, right))
|
|
if audio_lengths is None:
|
|
audio_lengths = torch.full((audio.shape[0],), original_length, device=audio.device, dtype=torch.long)
|
|
_, codes = self.quantizer(self.encoder(audio))
|
|
code_lengths = torch.ceil(audio_lengths.float() / self.frame_length).long()
|
|
max_codes = int(codes.shape[-1])
|
|
code_lengths = code_lengths.clamp_max(max_codes)
|
|
frame_index = torch.arange(max_codes, device=codes.device)
|
|
valid = frame_index[None, None, :] < code_lengths[:, None, None]
|
|
return torch.where(valid, codes, torch.full_like(codes, -1)), code_lengths
|
|
|
|
@torch.inference_mode()
|
|
def decode(self, codes: Tensor) -> Tensor:
|
|
"""Decode ``[B, 10, frames]`` to ``[B, 1, samples]``."""
|
|
return self.decoder(self.quantizer.decode(codes.long()))
|
|
|
|
|
|
def build_arktts_codec(
|
|
*,
|
|
post_n_layer: int = 8,
|
|
post_n_head: int = 16,
|
|
post_n_local_heads: int = 8,
|
|
post_intermediate_size: int = 1216,
|
|
) -> ArkttsCodec:
|
|
"""Construct the codec with uninitialised weights, on CPU, in eval mode."""
|
|
codec = ArkttsCodec(
|
|
post_n_layer=post_n_layer,
|
|
post_n_head=post_n_head,
|
|
post_n_local_heads=post_n_local_heads,
|
|
post_intermediate_size=post_intermediate_size,
|
|
)
|
|
codec.eval()
|
|
return codec
|
|
|
|
|
|
__all__ = ["ArkttsCodec", "build_arktts_codec"]
|