Sy03andYueqian Lin
58007c0084
Signed-off-by: Sysy <1370724210@qq.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
92 lines
2.8 KiB
Python
92 lines
2.8 KiB
Python
# SPDX-License-Identifier: MIT
|
|
# SPDX-FileCopyrightText: Copyright contributors to the Amphion project
|
|
"""Decode-only factorized residual vector quantizer.
|
|
|
|
Adapted from the Amphion codec implementation vendored by IndexTTS 2.5.
|
|
Training-only paths are intentionally omitted.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch.nn.utils import weight_norm
|
|
|
|
|
|
def _weight_norm_conv1d(*args, **kwargs) -> nn.Conv1d:
|
|
return weight_norm(nn.Conv1d(*args, **kwargs))
|
|
|
|
|
|
class FactorizedVectorQuantizeDecoder(nn.Module):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
input_dim: int,
|
|
codebook_size: int,
|
|
codebook_dim: int,
|
|
):
|
|
super().__init__()
|
|
self.input_dim = input_dim
|
|
self.codebook_dim = codebook_dim
|
|
if input_dim != codebook_dim:
|
|
# Keep upstream module names and weight-normalized parameter layout
|
|
# so codec.pth keys load without conversion.
|
|
self.in_project = _weight_norm_conv1d(
|
|
input_dim,
|
|
codebook_dim,
|
|
kernel_size=1,
|
|
)
|
|
self.out_project = _weight_norm_conv1d(
|
|
codebook_dim,
|
|
input_dim,
|
|
kernel_size=1,
|
|
)
|
|
else:
|
|
self.in_project = nn.Identity()
|
|
self.out_project = nn.Identity()
|
|
self.codebook = nn.Embedding(codebook_size, codebook_dim)
|
|
|
|
def vq2emb(self, codes: torch.Tensor) -> torch.Tensor:
|
|
embedding = F.embedding(codes, self.codebook.weight).transpose(1, 2)
|
|
return self.out_project(embedding)
|
|
|
|
|
|
class ResidualVQDecoder(nn.Module):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
input_dim: int,
|
|
num_quantizers: int,
|
|
codebook_size: int,
|
|
codebook_dim: int,
|
|
):
|
|
super().__init__()
|
|
self.num_quantizers = num_quantizers
|
|
self.quantizers = nn.ModuleList(
|
|
[
|
|
FactorizedVectorQuantizeDecoder(
|
|
input_dim=input_dim,
|
|
codebook_size=codebook_size,
|
|
codebook_dim=codebook_dim,
|
|
)
|
|
for _ in range(num_quantizers)
|
|
]
|
|
)
|
|
|
|
def vq2emb(
|
|
self,
|
|
codes: torch.Tensor,
|
|
n_quantizers: int | None = None,
|
|
) -> torch.Tensor:
|
|
count = self.num_quantizers if n_quantizers is None else n_quantizers
|
|
quantized: torch.Tensor | None = None
|
|
for index, quantizer in enumerate(self.quantizers):
|
|
if index >= count:
|
|
break
|
|
value = quantizer.vq2emb(codes[index])
|
|
quantized = value if quantized is None else quantized + value
|
|
if quantized is None:
|
|
raise ValueError("EnhancedCodec requires at least one quantizer")
|
|
return quantized
|