[Model] Add IndexTTS2 text-to-speech support (#3838)
Signed-off-by: BeatSeat <wendavid552@gmail.com> Signed-off-by: David <david@podcast.ai> Signed-off-by: wendavid552 <wendavid552@gmail.com> Signed-off-by: WenDavid <davidzhangcf@foxmail.com>
This commit is contained in:
BeatSeat authored and GitHub committed
1 parent
f4d922d8f6
commit
044240d8eb
66 files changed
+8634
-2
No files matched your search
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -0,0 +1,117 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
from collections import OrderedDict
|
||||
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from .layers import (
|
||||
BasicResBlock,
|
||||
CAMDenseTDNNBlock,
|
||||
DenseLayer,
|
||||
StatsPool,
|
||||
TDNNLayer,
|
||||
TransitLayer,
|
||||
get_nonlinear,
|
||||
)
|
||||
|
||||
|
||||
class FCM(nn.Module):
|
||||
def __init__(self, block=BasicResBlock, num_blocks=[2, 2], m_channels=32, feat_dim=80):
|
||||
super().__init__()
|
||||
self.in_planes = m_channels
|
||||
self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(m_channels)
|
||||
|
||||
self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=2)
|
||||
self.layer2 = self._make_layer(block, m_channels, num_blocks[1], stride=2)
|
||||
|
||||
self.conv2 = nn.Conv2d(m_channels, m_channels, kernel_size=3, stride=(2, 1), padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(m_channels)
|
||||
self.out_channels = m_channels * (feat_dim // 8)
|
||||
|
||||
def _make_layer(self, block, planes, num_blocks, stride):
|
||||
strides = [stride] + [1] * (num_blocks - 1)
|
||||
layers = []
|
||||
for stride in strides:
|
||||
layers.append(block(self.in_planes, planes, stride))
|
||||
self.in_planes = planes * block.expansion
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = x.unsqueeze(1)
|
||||
out = F.relu(self.bn1(self.conv1(x)))
|
||||
out = self.layer1(out)
|
||||
out = self.layer2(out)
|
||||
out = F.relu(self.bn2(self.conv2(out)))
|
||||
|
||||
shape = out.shape
|
||||
out = out.reshape(shape[0], shape[1] * shape[2], shape[3])
|
||||
return out
|
||||
|
||||
|
||||
class CAMPPlus(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
feat_dim=80,
|
||||
embedding_size=512,
|
||||
growth_rate=32,
|
||||
bn_size=4,
|
||||
init_channels=128,
|
||||
config_str="batchnorm-relu",
|
||||
memory_efficient=True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.head = FCM(feat_dim=feat_dim)
|
||||
channels = self.head.out_channels
|
||||
|
||||
self.xvector = nn.Sequential(
|
||||
OrderedDict(
|
||||
[
|
||||
(
|
||||
"tdnn",
|
||||
TDNNLayer(channels, init_channels, 5, stride=2, dilation=1, padding=-1, config_str=config_str),
|
||||
),
|
||||
]
|
||||
)
|
||||
)
|
||||
channels = init_channels
|
||||
for i, (num_layers, kernel_size, dilation) in enumerate(zip((12, 24, 16), (3, 3, 3), (1, 2, 2))):
|
||||
block = CAMDenseTDNNBlock(
|
||||
num_layers=num_layers,
|
||||
in_channels=channels,
|
||||
out_channels=growth_rate,
|
||||
bn_channels=bn_size * growth_rate,
|
||||
kernel_size=kernel_size,
|
||||
dilation=dilation,
|
||||
config_str=config_str,
|
||||
memory_efficient=memory_efficient,
|
||||
)
|
||||
self.xvector.add_module(f"block{i + 1}", block)
|
||||
channels = channels + num_layers * growth_rate
|
||||
self.xvector.add_module(
|
||||
f"transit{i + 1}", TransitLayer(channels, channels // 2, bias=False, config_str=config_str)
|
||||
)
|
||||
channels //= 2
|
||||
|
||||
self.xvector.add_module("out_nonlinear", get_nonlinear(config_str, channels))
|
||||
|
||||
self.xvector.add_module("stats", StatsPool())
|
||||
self.xvector.add_module("dense", DenseLayer(channels * 2, embedding_size, config_str="batchnorm_"))
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, (nn.Conv1d, nn.Linear)):
|
||||
nn.init.kaiming_normal_(m.weight.data)
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
def forward(self, x):
|
||||
x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
|
||||
x = self.head(x)
|
||||
x = self.xvector(x)
|
||||
return x
|
||||
@@ -0,0 +1,221 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.checkpoint as cp
|
||||
from torch import nn
|
||||
|
||||
|
||||
def get_nonlinear(config_str, channels):
|
||||
nonlinear = nn.Sequential()
|
||||
for name in config_str.split("-"):
|
||||
if name == "relu":
|
||||
nonlinear.add_module("relu", nn.ReLU(inplace=True))
|
||||
elif name == "prelu":
|
||||
nonlinear.add_module("prelu", nn.PReLU(channels))
|
||||
elif name == "batchnorm":
|
||||
nonlinear.add_module("batchnorm", nn.BatchNorm1d(channels))
|
||||
elif name == "batchnorm_":
|
||||
nonlinear.add_module("batchnorm", nn.BatchNorm1d(channels, affine=False))
|
||||
else:
|
||||
raise ValueError(f"Unexpected module ({name}).")
|
||||
return nonlinear
|
||||
|
||||
|
||||
def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, eps=1e-2):
|
||||
mean = x.mean(dim=dim)
|
||||
std = x.std(dim=dim, unbiased=unbiased)
|
||||
stats = torch.cat([mean, std], dim=-1)
|
||||
if keepdim:
|
||||
stats = stats.unsqueeze(dim=dim)
|
||||
return stats
|
||||
|
||||
|
||||
class StatsPool(nn.Module):
|
||||
def forward(self, x):
|
||||
return statistics_pooling(x)
|
||||
|
||||
|
||||
class TDNNLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
bias=False,
|
||||
config_str="batchnorm-relu",
|
||||
):
|
||||
super().__init__()
|
||||
if padding < 0:
|
||||
assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})"
|
||||
padding = (kernel_size - 1) // 2 * dilation
|
||||
self.linear = nn.Conv1d(
|
||||
in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias
|
||||
)
|
||||
self.nonlinear = get_nonlinear(config_str, out_channels)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.linear(x)
|
||||
x = self.nonlinear(x)
|
||||
return x
|
||||
|
||||
|
||||
class CAMLayer(nn.Module):
|
||||
def __init__(self, bn_channels, out_channels, kernel_size, stride, padding, dilation, bias, reduction=2):
|
||||
super().__init__()
|
||||
self.linear_local = nn.Conv1d(
|
||||
bn_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias
|
||||
)
|
||||
self.linear1 = nn.Conv1d(bn_channels, bn_channels // reduction, 1)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.linear2 = nn.Conv1d(bn_channels // reduction, out_channels, 1)
|
||||
self.sigmoid = nn.Sigmoid()
|
||||
|
||||
def forward(self, x):
|
||||
y = self.linear_local(x)
|
||||
context = x.mean(-1, keepdim=True) + self.seg_pooling(x)
|
||||
context = self.relu(self.linear1(context))
|
||||
m = self.sigmoid(self.linear2(context))
|
||||
return y * m
|
||||
|
||||
def seg_pooling(self, x, seg_len=100, stype="avg"):
|
||||
if stype == "avg":
|
||||
seg = F.avg_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
|
||||
elif stype == "max":
|
||||
seg = F.max_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
|
||||
else:
|
||||
raise ValueError("Wrong segment pooling type.")
|
||||
shape = seg.shape
|
||||
seg = seg.unsqueeze(-1).expand(*shape, seg_len).reshape(*shape[:-1], -1)
|
||||
seg = seg[..., : x.shape[-1]]
|
||||
return seg
|
||||
|
||||
|
||||
class CAMDenseTDNNLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
bn_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
dilation=1,
|
||||
bias=False,
|
||||
config_str="batchnorm-relu",
|
||||
memory_efficient=False,
|
||||
):
|
||||
super().__init__()
|
||||
assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})"
|
||||
padding = (kernel_size - 1) // 2 * dilation
|
||||
self.memory_efficient = memory_efficient
|
||||
self.nonlinear1 = get_nonlinear(config_str, in_channels)
|
||||
self.linear1 = nn.Conv1d(in_channels, bn_channels, 1, bias=False)
|
||||
self.nonlinear2 = get_nonlinear(config_str, bn_channels)
|
||||
self.cam_layer = CAMLayer(
|
||||
bn_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias
|
||||
)
|
||||
|
||||
def bn_function(self, x):
|
||||
return self.linear1(self.nonlinear1(x))
|
||||
|
||||
def forward(self, x):
|
||||
if self.training and self.memory_efficient:
|
||||
x = cp.checkpoint(self.bn_function, x)
|
||||
else:
|
||||
x = self.bn_function(x)
|
||||
x = self.cam_layer(self.nonlinear2(x))
|
||||
return x
|
||||
|
||||
|
||||
class CAMDenseTDNNBlock(nn.ModuleList):
|
||||
def __init__(
|
||||
self,
|
||||
num_layers,
|
||||
in_channels,
|
||||
out_channels,
|
||||
bn_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
dilation=1,
|
||||
bias=False,
|
||||
config_str="batchnorm-relu",
|
||||
memory_efficient=False,
|
||||
):
|
||||
super().__init__()
|
||||
for i in range(num_layers):
|
||||
layer = CAMDenseTDNNLayer(
|
||||
in_channels=in_channels + i * out_channels,
|
||||
out_channels=out_channels,
|
||||
bn_channels=bn_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
dilation=dilation,
|
||||
bias=bias,
|
||||
config_str=config_str,
|
||||
memory_efficient=memory_efficient,
|
||||
)
|
||||
self.add_module(f"tdnnd{i + 1}", layer)
|
||||
|
||||
def forward(self, x):
|
||||
for layer in self:
|
||||
x = torch.cat([x, layer(x)], dim=1)
|
||||
return x
|
||||
|
||||
|
||||
class TransitLayer(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, bias=True, config_str="batchnorm-relu"):
|
||||
super().__init__()
|
||||
self.nonlinear = get_nonlinear(config_str, in_channels)
|
||||
self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.nonlinear(x)
|
||||
x = self.linear(x)
|
||||
return x
|
||||
|
||||
|
||||
class DenseLayer(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, bias=False, config_str="batchnorm-relu"):
|
||||
super().__init__()
|
||||
self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias)
|
||||
self.nonlinear = get_nonlinear(config_str, out_channels)
|
||||
|
||||
def forward(self, x):
|
||||
if len(x.shape) == 2:
|
||||
x = self.linear(x.unsqueeze(dim=-1)).squeeze(dim=-1)
|
||||
else:
|
||||
x = self.linear(x)
|
||||
x = self.nonlinear(x)
|
||||
return x
|
||||
|
||||
|
||||
class BasicResBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, in_planes, planes, stride=1):
|
||||
super().__init__()
|
||||
self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=(stride, 1), padding=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(planes)
|
||||
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(planes)
|
||||
|
||||
self.shortcut = nn.Sequential()
|
||||
if stride != 1 or in_planes != self.expansion * planes:
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=(stride, 1), bias=False),
|
||||
nn.BatchNorm2d(self.expansion * planes),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
out = F.relu(self.bn1(self.conv1(x)))
|
||||
out = self.bn2(self.conv2(out))
|
||||
out += self.shortcut(x)
|
||||
out = F.relu(out)
|
||||
return out
|
||||
@@ -0,0 +1,88 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import re
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def tokenize_by_CJK_char(line: str, do_upper_case=True) -> str:
|
||||
"""
|
||||
Tokenize a line of text with CJK char.
|
||||
|
||||
Note: All return characters will be upper case.
|
||||
|
||||
Example:
|
||||
input = "你好世界是 hello world 的中文"
|
||||
output = "你 好 世 界 是 HELLO WORLD 的 中 文"
|
||||
|
||||
Args:
|
||||
line:
|
||||
The input text.
|
||||
|
||||
Return:
|
||||
A new string tokenize by CJK char.
|
||||
"""
|
||||
# CJK ranges from nltk/tokenize/util.py (alvations/nltk@79eed6d)
|
||||
CJK_RANGE_PATTERN = (
|
||||
r"([\u1100-\u11ff\u2e80-\ua4cf\ua840-\uD7AF\uF900-\uFAFF\uFE30-\uFE4F\uFF65-\uFFDC\U00020000-\U0002FFFF])"
|
||||
)
|
||||
chars = re.split(CJK_RANGE_PATTERN, line.strip())
|
||||
return " ".join([w.strip().upper() if do_upper_case else w.strip() for w in chars if w.strip()])
|
||||
|
||||
|
||||
def de_tokenized_by_CJK_char(line: str, do_lower_case=False) -> str:
|
||||
"""
|
||||
Example:
|
||||
input = "你 好 世 界 是 HELLO WORLD 的 中 文"
|
||||
output = "你好世界是 hello world 的中文"
|
||||
|
||||
do_lower_case:
|
||||
input = "SEE YOU!"
|
||||
output = "see you!"
|
||||
"""
|
||||
# replace english words in the line with placeholders
|
||||
english_word_pattern = re.compile(r"([A-Z]+(?:[\s'-][A-Z-]+)*)", re.IGNORECASE)
|
||||
english_sents = english_word_pattern.findall(line)
|
||||
for i, sent in enumerate(english_sents):
|
||||
line = line.replace(sent, f"<sent_{i}>")
|
||||
|
||||
words = line.split()
|
||||
# restore english sentences
|
||||
sent_placeholder_pattern = re.compile(r"(<sent_(\d+)>)")
|
||||
for i in range(len(words)):
|
||||
all_matches = sent_placeholder_pattern.findall(words[i])
|
||||
if all_matches:
|
||||
# restore the english word
|
||||
for h, j in all_matches:
|
||||
placeholder_index = int(j)
|
||||
words[i] = words[i].replace(h, english_sents[placeholder_index])
|
||||
if do_lower_case:
|
||||
words[i] = words[i].lower()
|
||||
return "".join(words)
|
||||
|
||||
|
||||
def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
|
||||
"""Make mask tensor containing indices of padded part.
|
||||
|
||||
See description of make_non_pad_mask.
|
||||
|
||||
Args:
|
||||
lengths (torch.Tensor): Batch of lengths (B,).
|
||||
Returns:
|
||||
torch.Tensor: Mask tensor containing indices of padded part.
|
||||
|
||||
Examples:
|
||||
>>> lengths = [5, 3, 2]
|
||||
>>> make_pad_mask(lengths)
|
||||
masks = [[0, 0, 0, 0 ,0],
|
||||
[0, 0, 0, 1, 1],
|
||||
[0, 0, 1, 1, 1]]
|
||||
"""
|
||||
batch_size = lengths.size(0)
|
||||
max_len = max_len if max_len > 0 else lengths.max().item()
|
||||
seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device)
|
||||
seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
|
||||
seq_length_expand = lengths.unsqueeze(-1)
|
||||
mask = seq_range_expand >= seq_length_expand
|
||||
return mask
|
||||
@@ -0,0 +1,208 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import os
|
||||
import warnings
|
||||
from typing import overload
|
||||
|
||||
from sentencepiece import SentencePieceProcessor
|
||||
|
||||
from vllm_omni.model_executor.models.indextts2.utils.common import (
|
||||
de_tokenized_by_CJK_char,
|
||||
tokenize_by_CJK_char,
|
||||
)
|
||||
|
||||
|
||||
class TextTokenizer:
|
||||
def __init__(self, vocab_file: str):
|
||||
self.vocab_file = vocab_file
|
||||
|
||||
if self.vocab_file is None:
|
||||
raise ValueError("vocab_file is None")
|
||||
if not os.path.exists(self.vocab_file):
|
||||
raise ValueError(f"vocab_file {self.vocab_file} does not exist")
|
||||
self.sp_model = SentencePieceProcessor(model_file=self.vocab_file)
|
||||
self.pre_tokenizers = [tokenize_by_CJK_char]
|
||||
|
||||
@property
|
||||
def vocab_size(self):
|
||||
return self.sp_model.GetPieceSize()
|
||||
|
||||
@property
|
||||
def unk_token(self):
|
||||
return "<unk>"
|
||||
|
||||
@property
|
||||
def pad_token(self):
|
||||
return None
|
||||
|
||||
@property
|
||||
def bos_token(self):
|
||||
return "<s>"
|
||||
|
||||
@property
|
||||
def eos_token(self):
|
||||
return "</s>"
|
||||
|
||||
@property
|
||||
def pad_token_id(self):
|
||||
return -1
|
||||
|
||||
@property
|
||||
def bos_token_id(self):
|
||||
return 0
|
||||
|
||||
@property
|
||||
def eos_token_id(self):
|
||||
return 1
|
||||
|
||||
@property
|
||||
def unk_token_id(self):
|
||||
return self.sp_model.unk_id()
|
||||
|
||||
@property
|
||||
def special_tokens_map(self):
|
||||
return {
|
||||
"unk_token": self.unk_token,
|
||||
"pad_token": self.pad_token,
|
||||
"bos_token": self.bos_token,
|
||||
"eos_token": self.eos_token,
|
||||
}
|
||||
|
||||
def get_vocab(self):
|
||||
vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
|
||||
return vocab
|
||||
|
||||
@overload
|
||||
def convert_ids_to_tokens(self, ids: int) -> str: ...
|
||||
|
||||
@overload
|
||||
def convert_ids_to_tokens(self, ids: list[int]) -> list[str]: ...
|
||||
|
||||
def convert_ids_to_tokens(self, ids: list[int] | int):
|
||||
return self.sp_model.IdToPiece(ids)
|
||||
|
||||
def convert_tokens_to_ids(self, tokens: list[str] | str) -> list[int]:
|
||||
if isinstance(tokens, str):
|
||||
tokens = [tokens]
|
||||
return [self.sp_model.PieceToId(token) for token in tokens]
|
||||
|
||||
def tokenize(self, text: str) -> list[str]:
|
||||
return self.encode(text, out_type=str)
|
||||
|
||||
def encode(self, text: str, **kwargs):
|
||||
if len(text) == 0:
|
||||
return []
|
||||
if len(text.strip()) == 1:
|
||||
return self.sp_model.Encode(text, out_type=kwargs.pop("out_type", int), **kwargs)
|
||||
for pre_tokenizer in self.pre_tokenizers:
|
||||
text = pre_tokenizer(text)
|
||||
return self.sp_model.Encode(text, out_type=kwargs.pop("out_type", int), **kwargs)
|
||||
|
||||
def decode(self, ids: list[int] | int, do_lower_case=False, **kwargs):
|
||||
if isinstance(ids, int):
|
||||
ids = [ids]
|
||||
decoded = self.sp_model.Decode(ids, out_type=kwargs.pop("out_type", str), **kwargs)
|
||||
return de_tokenized_by_CJK_char(decoded, do_lower_case=do_lower_case)
|
||||
|
||||
@staticmethod
|
||||
def split_segments_by_token(
|
||||
tokenized_str: list[str],
|
||||
split_tokens: list[str],
|
||||
max_text_tokens_per_segment: int,
|
||||
quick_streaming_tokens: int = 0,
|
||||
) -> list[list[str]]:
|
||||
if len(tokenized_str) == 0:
|
||||
return []
|
||||
segments: list[list[str]] = []
|
||||
current_segment = []
|
||||
current_segment_tokens_len = 0
|
||||
i = 0
|
||||
while i < len(tokenized_str):
|
||||
token = tokenized_str[i]
|
||||
current_segment.append(token)
|
||||
current_segment_tokens_len += 1
|
||||
if not ("," in split_tokens or "▁," in split_tokens) and (
|
||||
"," in current_segment or "▁," in current_segment
|
||||
):
|
||||
sub_segments = TextTokenizer.split_segments_by_token(
|
||||
current_segment,
|
||||
[",", "▁,"],
|
||||
max_text_tokens_per_segment=max_text_tokens_per_segment,
|
||||
quick_streaming_tokens=quick_streaming_tokens,
|
||||
)
|
||||
elif "-" not in split_tokens and "-" in current_segment:
|
||||
sub_segments = TextTokenizer.split_segments_by_token(
|
||||
current_segment,
|
||||
["-"],
|
||||
max_text_tokens_per_segment=max_text_tokens_per_segment,
|
||||
quick_streaming_tokens=quick_streaming_tokens,
|
||||
)
|
||||
elif current_segment_tokens_len <= max_text_tokens_per_segment:
|
||||
if token in split_tokens and current_segment_tokens_len > 2:
|
||||
if i < len(tokenized_str) - 1:
|
||||
if tokenized_str[i + 1] in ["'", "▁'"]:
|
||||
current_segment.append(tokenized_str[i + 1])
|
||||
i += 1
|
||||
segments.append(current_segment)
|
||||
current_segment = []
|
||||
current_segment_tokens_len = 0
|
||||
i += 1
|
||||
continue
|
||||
else:
|
||||
sub_segments = []
|
||||
for j in range(0, len(current_segment), max_text_tokens_per_segment):
|
||||
if j + max_text_tokens_per_segment < len(current_segment):
|
||||
sub_segments.append(current_segment[j : j + max_text_tokens_per_segment])
|
||||
else:
|
||||
sub_segments.append(current_segment[j:])
|
||||
warnings.warn(
|
||||
f"The tokens length of segment exceeds limit: {max_text_tokens_per_segment}, "
|
||||
f"Tokens in segment: {current_segment}."
|
||||
"Maybe unexpected behavior",
|
||||
RuntimeWarning,
|
||||
)
|
||||
segments.extend(sub_segments)
|
||||
current_segment = []
|
||||
current_segment_tokens_len = 0
|
||||
i += 1
|
||||
if current_segment_tokens_len > 0:
|
||||
assert current_segment_tokens_len <= max_text_tokens_per_segment
|
||||
segments.append(current_segment)
|
||||
merged_segments = []
|
||||
total_token = 0
|
||||
for segment in segments:
|
||||
total_token += len(segment)
|
||||
if len(segment) == 0:
|
||||
continue
|
||||
if len(merged_segments) == 0:
|
||||
merged_segments.append(segment)
|
||||
elif (
|
||||
len(merged_segments[-1]) + len(segment) <= max_text_tokens_per_segment
|
||||
and total_token > quick_streaming_tokens
|
||||
):
|
||||
merged_segments[-1] = merged_segments[-1] + segment
|
||||
elif len(merged_segments[-1]) + len(segment) <= max_text_tokens_per_segment / 2:
|
||||
merged_segments[-1] = merged_segments[-1] + segment
|
||||
else:
|
||||
merged_segments.append(segment)
|
||||
return merged_segments
|
||||
|
||||
punctuation_marks_tokens = [
|
||||
".",
|
||||
"!",
|
||||
"?",
|
||||
"▁.",
|
||||
"▁?",
|
||||
"▁...",
|
||||
]
|
||||
|
||||
def split_segments(
|
||||
self, tokenized: list[str], max_text_tokens_per_segment=120, quick_streaming_tokens=0
|
||||
) -> list[list[str]]:
|
||||
return TextTokenizer.split_segments_by_token(
|
||||
tokenized,
|
||||
self.punctuation_marks_tokens,
|
||||
max_text_tokens_per_segment=max_text_tokens_per_segment,
|
||||
quick_streaming_tokens=quick_streaming_tokens,
|
||||
)
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -0,0 +1,200 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Copyright (c) 2024 Amphion.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
#
|
||||
# Merged from factorized_vector_quantize.py and residual_vq.py.
|
||||
# Only FactorizedVectorQuantize ("fvq") quantizer type is retained
|
||||
# as it is the only one used by RepCodec in IndexTTS2.
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
from torch.nn.utils import weight_norm
|
||||
|
||||
|
||||
def WNConv1d(*args, **kwargs):
|
||||
return weight_norm(nn.Conv1d(*args, **kwargs))
|
||||
|
||||
|
||||
class FactorizedVectorQuantize(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_dim,
|
||||
codebook_size,
|
||||
codebook_dim,
|
||||
commitment=0.005,
|
||||
codebook_loss_weight=1.0,
|
||||
use_l2_normalize=True,
|
||||
):
|
||||
super().__init__()
|
||||
self.input_dim = input_dim
|
||||
self.codebook_size = codebook_size
|
||||
self.codebook_dim = codebook_dim
|
||||
self.commitment = commitment
|
||||
self.codebook_loss_weight = codebook_loss_weight
|
||||
self.use_l2_normalize = use_l2_normalize
|
||||
|
||||
if self.input_dim != self.codebook_dim:
|
||||
self.in_project = WNConv1d(self.input_dim, self.codebook_dim, kernel_size=1)
|
||||
self.out_project = WNConv1d(self.codebook_dim, self.input_dim, kernel_size=1)
|
||||
else:
|
||||
self.in_project = nn.Identity()
|
||||
self.out_project = nn.Identity()
|
||||
|
||||
self.codebook = nn.Embedding(self.codebook_size, self.codebook_dim)
|
||||
|
||||
def forward(self, z):
|
||||
z_e = self.in_project(z)
|
||||
z_q, indices = self.decode_latents(z_e)
|
||||
|
||||
if self.training:
|
||||
commit_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2]) * self.commitment
|
||||
codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2]) * self.codebook_loss_weight
|
||||
else:
|
||||
commit_loss = torch.zeros(z.shape[0], device=z.device)
|
||||
codebook_loss = torch.zeros(z.shape[0], device=z.device)
|
||||
|
||||
z_q = z_e + (z_q - z_e).detach()
|
||||
z_q = self.out_project(z_q)
|
||||
|
||||
return z_q, commit_loss, codebook_loss, indices, z_e
|
||||
|
||||
def embed_code(self, embed_id):
|
||||
return F.embedding(embed_id, self.codebook.weight)
|
||||
|
||||
def decode_code(self, embed_id):
|
||||
return self.embed_code(embed_id).transpose(1, 2)
|
||||
|
||||
def decode_latents(self, latents):
|
||||
encodings = rearrange(latents, "b d t -> (b t) d")
|
||||
codebook = self.codebook.weight
|
||||
|
||||
if self.use_l2_normalize:
|
||||
encodings = F.normalize(encodings)
|
||||
codebook = F.normalize(codebook)
|
||||
|
||||
dist = (
|
||||
encodings.pow(2).sum(1, keepdim=True)
|
||||
- 2 * encodings @ codebook.t()
|
||||
+ codebook.pow(2).sum(1, keepdim=True).t()
|
||||
)
|
||||
indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
|
||||
z_q = self.decode_code(indices)
|
||||
return z_q, indices
|
||||
|
||||
def vq2emb(self, vq, out_proj=True):
|
||||
emb = self.decode_code(vq)
|
||||
if out_proj:
|
||||
emb = self.out_project(emb)
|
||||
return emb
|
||||
|
||||
|
||||
class ResidualVQ(nn.Module):
|
||||
"""Residual Vector Quantization (SoundStream).
|
||||
|
||||
Only supports FactorizedVectorQuantize ("fvq") quantizer type.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_dim: int = 256,
|
||||
num_quantizers: int = 8,
|
||||
codebook_size: int = 1024,
|
||||
codebook_dim: int = 256,
|
||||
quantizer_type: str = "fvq",
|
||||
quantizer_dropout: float = 0.5,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
self.input_dim = input_dim
|
||||
self.num_quantizers = num_quantizers
|
||||
self.codebook_size = codebook_size
|
||||
self.codebook_dim = codebook_dim
|
||||
self.quantizer_dropout = quantizer_dropout
|
||||
|
||||
self.quantizers = nn.ModuleList(
|
||||
[
|
||||
FactorizedVectorQuantize(
|
||||
input_dim=input_dim,
|
||||
codebook_size=codebook_size,
|
||||
codebook_dim=codebook_dim,
|
||||
**kwargs,
|
||||
)
|
||||
for _ in range(num_quantizers)
|
||||
]
|
||||
)
|
||||
|
||||
def forward(self, z, n_quantizers: int = None):
|
||||
quantized_out = 0.0
|
||||
residual = z
|
||||
|
||||
all_commit_losses = []
|
||||
all_codebook_losses = []
|
||||
all_indices = []
|
||||
all_quantized = []
|
||||
|
||||
if n_quantizers is None:
|
||||
n_quantizers = self.num_quantizers
|
||||
|
||||
if self.training:
|
||||
n_q = torch.ones((z.shape[0],)) * self.num_quantizers + 1
|
||||
dropout = torch.randint(1, self.num_quantizers + 1, (z.shape[0],))
|
||||
n_dropout = int(z.shape[0] * self.quantizer_dropout)
|
||||
n_q[:n_dropout] = dropout[:n_dropout]
|
||||
n_quantizers = n_q.to(z.device)
|
||||
|
||||
for i, quantizer in enumerate(self.quantizers):
|
||||
if not self.training and i >= n_quantizers:
|
||||
break
|
||||
|
||||
z_q_i, commit_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(residual)
|
||||
|
||||
mask = torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers
|
||||
quantized_out = quantized_out + z_q_i * mask[:, None, None]
|
||||
residual = residual - z_q_i
|
||||
|
||||
commit_loss_i = (commit_loss_i * mask).mean()
|
||||
codebook_loss_i = (codebook_loss_i * mask).mean()
|
||||
|
||||
all_commit_losses.append(commit_loss_i)
|
||||
all_codebook_losses.append(codebook_loss_i)
|
||||
all_indices.append(indices_i)
|
||||
all_quantized.append(z_q_i)
|
||||
|
||||
(
|
||||
all_commit_losses,
|
||||
all_codebook_losses,
|
||||
all_indices,
|
||||
all_quantized,
|
||||
) = map(
|
||||
torch.stack,
|
||||
(
|
||||
all_commit_losses,
|
||||
all_codebook_losses,
|
||||
all_indices,
|
||||
all_quantized,
|
||||
),
|
||||
)
|
||||
|
||||
return (
|
||||
quantized_out,
|
||||
all_indices,
|
||||
all_commit_losses,
|
||||
all_codebook_losses,
|
||||
all_quantized,
|
||||
)
|
||||
|
||||
def vq2emb(self, vq, n_quantizers=None):
|
||||
quantized_out = 0.0
|
||||
if n_quantizers is None:
|
||||
n_quantizers = self.num_quantizers
|
||||
for idx, quantizer in enumerate(self.quantizers):
|
||||
if idx >= n_quantizers:
|
||||
break
|
||||
quantized_out += quantizer.vq2emb(vq[idx])
|
||||
return quantized_out
|
||||
@@ -0,0 +1,120 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Copyright (c) 2024 Amphion.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch.nn as nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from vllm_omni.model_executor.models.indextts2.utils.maskgct.quantize import ResidualVQ
|
||||
from vllm_omni.model_executor.models.indextts2.utils.maskgct.vocos import VocosBackbone
|
||||
|
||||
|
||||
def init_weights(m):
|
||||
if isinstance(m, nn.Conv1d):
|
||||
nn.init.trunc_normal_(m.weight, std=0.02)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.trunc_normal_(m.weight, std=0.02)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
|
||||
class RepCodec(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
codebook_size=8192,
|
||||
hidden_size=1024,
|
||||
codebook_dim=8,
|
||||
vocos_dim=384,
|
||||
vocos_intermediate_dim=2048,
|
||||
vocos_num_layers=12,
|
||||
num_quantizers=1,
|
||||
downsample_scale=1,
|
||||
cfg=None,
|
||||
):
|
||||
super().__init__()
|
||||
if cfg is not None:
|
||||
_g = lambda attr, default: getattr(cfg, attr, default) # noqa: E731
|
||||
codebook_size = _g("codebook_size", codebook_size)
|
||||
codebook_dim = _g("codebook_dim", codebook_dim)
|
||||
hidden_size = _g("hidden_size", hidden_size)
|
||||
vocos_dim = _g("vocos_dim", vocos_dim)
|
||||
vocos_intermediate_dim = _g("vocos_intermediate_dim", vocos_intermediate_dim)
|
||||
vocos_num_layers = _g("vocos_num_layers", vocos_num_layers)
|
||||
num_quantizers = _g("num_quantizers", num_quantizers)
|
||||
downsample_scale = _g("downsample_scale", downsample_scale)
|
||||
|
||||
self.codebook_size = codebook_size
|
||||
self.codebook_dim = codebook_dim
|
||||
self.hidden_size = hidden_size
|
||||
self.vocos_dim = vocos_dim
|
||||
self.vocos_intermediate_dim = vocos_intermediate_dim
|
||||
self.vocos_num_layers = vocos_num_layers
|
||||
self.num_quantizers = num_quantizers
|
||||
self.downsample_scale = downsample_scale
|
||||
|
||||
if self.downsample_scale is not None and self.downsample_scale > 1:
|
||||
self.down = nn.Conv1d(self.hidden_size, self.hidden_size, kernel_size=3, stride=2, padding=1)
|
||||
self.up = nn.Conv1d(self.hidden_size, self.hidden_size, kernel_size=3, stride=1, padding=1)
|
||||
|
||||
self.encoder = nn.Sequential(
|
||||
VocosBackbone(
|
||||
input_channels=self.hidden_size,
|
||||
dim=self.vocos_dim,
|
||||
intermediate_dim=self.vocos_intermediate_dim,
|
||||
num_layers=self.vocos_num_layers,
|
||||
adanorm_num_embeddings=None,
|
||||
),
|
||||
nn.Linear(self.vocos_dim, self.hidden_size),
|
||||
)
|
||||
self.decoder = nn.Sequential(
|
||||
VocosBackbone(
|
||||
input_channels=self.hidden_size,
|
||||
dim=self.vocos_dim,
|
||||
intermediate_dim=self.vocos_intermediate_dim,
|
||||
num_layers=self.vocos_num_layers,
|
||||
adanorm_num_embeddings=None,
|
||||
),
|
||||
nn.Linear(self.vocos_dim, self.hidden_size),
|
||||
)
|
||||
|
||||
self.quantizer = ResidualVQ(
|
||||
input_dim=hidden_size,
|
||||
num_quantizers=num_quantizers,
|
||||
codebook_size=codebook_size,
|
||||
codebook_dim=codebook_dim,
|
||||
quantizer_type="fvq",
|
||||
quantizer_dropout=0.0,
|
||||
commitment=0.15,
|
||||
codebook_loss_weight=1.0,
|
||||
use_l2_normalize=True,
|
||||
)
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def quantize(self, x):
|
||||
if self.downsample_scale is not None and self.downsample_scale > 1:
|
||||
x = x.transpose(1, 2)
|
||||
x = self.down(x)
|
||||
x = F.gelu(x)
|
||||
x = x.transpose(1, 2)
|
||||
|
||||
x = self.encoder(x.transpose(1, 2)).transpose(1, 2)
|
||||
|
||||
(
|
||||
quantized_out,
|
||||
all_indices,
|
||||
all_commit_losses,
|
||||
all_codebook_losses,
|
||||
_,
|
||||
) = self.quantizer(x)
|
||||
|
||||
if all_indices.shape[0] == 1:
|
||||
return all_indices.squeeze(0), quantized_out.transpose(1, 2)
|
||||
return all_indices, quantized_out.transpose(1, 2)
|
||||
|
||||
def reset_parameters(self):
|
||||
self.apply(init_weights)
|
||||
@@ -0,0 +1,102 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Copyright (c) 2024 Amphion.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class ConvNeXtBlock(nn.Module):
|
||||
"""ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
intermediate_dim: int,
|
||||
layer_scale_init_value: float,
|
||||
adanorm_num_embeddings: int | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim)
|
||||
self.adanorm = adanorm_num_embeddings is not None
|
||||
if adanorm_num_embeddings:
|
||||
raise ValueError("AdaLayerNorm is not supported in inference-only mode")
|
||||
else:
|
||||
self.norm = nn.LayerNorm(dim, eps=1e-6)
|
||||
self.pwconv1 = nn.Linear(dim, intermediate_dim)
|
||||
self.act = nn.GELU()
|
||||
self.pwconv2 = nn.Linear(intermediate_dim, dim)
|
||||
self.gamma = (
|
||||
nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
|
||||
if layer_scale_init_value > 0
|
||||
else None
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor | None = None) -> torch.Tensor:
|
||||
residual = x
|
||||
x = self.dwconv(x)
|
||||
x = x.transpose(1, 2)
|
||||
x = self.norm(x)
|
||||
x = self.pwconv1(x)
|
||||
x = self.act(x)
|
||||
x = self.pwconv2(x)
|
||||
if self.gamma is not None:
|
||||
x = self.gamma * x
|
||||
x = x.transpose(1, 2)
|
||||
x = residual + x
|
||||
return x
|
||||
|
||||
|
||||
class VocosBackbone(nn.Module):
|
||||
"""Vocos backbone module built with ConvNeXt blocks."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_channels: int,
|
||||
dim: int,
|
||||
intermediate_dim: int,
|
||||
num_layers: int,
|
||||
layer_scale_init_value: float | None = None,
|
||||
adanorm_num_embeddings: int | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.input_channels = input_channels
|
||||
self.embed = nn.Conv1d(input_channels, dim, kernel_size=7, padding=3)
|
||||
self.adanorm = adanorm_num_embeddings is not None
|
||||
if adanorm_num_embeddings:
|
||||
raise ValueError("AdaLayerNorm is not supported in inference-only mode")
|
||||
else:
|
||||
self.norm = nn.LayerNorm(dim, eps=1e-6)
|
||||
layer_scale_init_value = layer_scale_init_value or 1 / num_layers
|
||||
self.convnext = nn.ModuleList(
|
||||
[
|
||||
ConvNeXtBlock(
|
||||
dim=dim,
|
||||
intermediate_dim=intermediate_dim,
|
||||
layer_scale_init_value=layer_scale_init_value,
|
||||
adanorm_num_embeddings=adanorm_num_embeddings,
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
]
|
||||
)
|
||||
self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6)
|
||||
self.apply(self._init_weights)
|
||||
|
||||
def _init_weights(self, m):
|
||||
if isinstance(m, (nn.Conv1d, nn.Linear)):
|
||||
nn.init.trunc_normal_(m.weight, std=0.02)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
|
||||
x = self.embed(x)
|
||||
x = self.norm(x.transpose(1, 2))
|
||||
x = x.transpose(1, 2)
|
||||
for conv_block in self.convnext:
|
||||
x = conv_block(x)
|
||||
x = self.final_layer_norm(x.transpose(1, 2))
|
||||
return x
|
||||
Reference in new issue
Block a user