From 4cd1d83de1ff5e4339cb85448ed4e3e69066e3e5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CE=9D=CE=B1=CF=81=CE=BF=CF=85=CF=83=CE=AD=C2=B7=CE=BC?= =?UTF-8?q?=C2=B7=CE=B3=CE=B9=CE=BF=CF=85=CE=BC=CE=B5=CE=BC=CE=AF=C2=B7?= =?UTF-8?q?=CE=A7=CE=B9=CE=BD=CE=B1=CE=BA=CE=AC=CE=BD=CE=BD=CE=B1?= <40709280+NaruseMioShirakana@users.noreply.github.com> Date: Sat, 6 Apr 2024 15:45:51 +0800 Subject: [PATCH 01/14] Add files via upload --- GPT_SoVITS/module/activation_onnx.py | 178 +++++++++++ GPT_SoVITS/module/embedding_onnx.py | 64 ++++ .../module/patched_mha_with_cache_onnx.py | 92 ++++++ GPT_SoVITS/module/transformer_onnx.py | 292 ++++++++++++++++++ 4 files changed, 626 insertions(+) create mode 100644 GPT_SoVITS/module/activation_onnx.py create mode 100644 GPT_SoVITS/module/embedding_onnx.py create mode 100644 GPT_SoVITS/module/patched_mha_with_cache_onnx.py create mode 100644 GPT_SoVITS/module/transformer_onnx.py diff --git a/GPT_SoVITS/module/activation_onnx.py b/GPT_SoVITS/module/activation_onnx.py new file mode 100644 index 0000000..b54acd9 --- /dev/null +++ b/GPT_SoVITS/module/activation_onnx.py @@ -0,0 +1,178 @@ +# modified from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/activation.py +from typing import Optional +from typing import Tuple +import torch +from torch import Tensor +from torch.nn import Linear +from torch.nn import Module +from torch.nn.init import constant_ +from torch.nn.init import xavier_normal_ +from torch.nn.init import xavier_uniform_ +from torch.nn.modules.linear import NonDynamicallyQuantizableLinear +from torch.nn.parameter import Parameter + +from torch.nn import functional as F +from AR.modules.patched_mha_with_cache_onnx import multi_head_attention_forward_patched + + +class MultiheadAttention(Module): + __constants__ = ["batch_first"] + bias_k: Optional[torch.Tensor] + bias_v: Optional[torch.Tensor] + + def __init__( + self, + embed_dim, + num_heads, + dropout=0.0, + bias=True, + add_bias_kv=False, + add_zero_attn=False, + kdim=None, + vdim=None, + batch_first=False, + linear1_cls=Linear, + linear2_cls=Linear, + device=None, + dtype=None, + ) -> None: + factory_kwargs = {"device": device, "dtype": dtype} + super(MultiheadAttention, self).__init__() + self.embed_dim = embed_dim + self.kdim = kdim if kdim is not None else embed_dim + self.vdim = vdim if vdim is not None else embed_dim + self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim + + self.num_heads = num_heads + self.dropout = dropout + self.batch_first = batch_first + self.head_dim = embed_dim // num_heads + assert ( + self.head_dim * num_heads == self.embed_dim + ), "embed_dim must be divisible by num_heads" + + if add_bias_kv: + self.bias_k = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs)) + self.bias_v = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs)) + else: + self.bias_k = self.bias_v = None + + if linear1_cls == Linear: + if not self._qkv_same_embed_dim: + self.q_proj_weight = Parameter( + torch.empty((embed_dim, embed_dim), **factory_kwargs) + ) + self.k_proj_weight = Parameter( + torch.empty((embed_dim, self.kdim), **factory_kwargs) + ) + self.v_proj_weight = Parameter( + torch.empty((embed_dim, self.vdim), **factory_kwargs) + ) + self.register_parameter("in_proj_weight", None) + else: + self.in_proj_weight = Parameter( + torch.empty((3 * embed_dim, embed_dim), **factory_kwargs) + ) + self.register_parameter("q_proj_weight", None) + self.register_parameter("k_proj_weight", None) + self.register_parameter("v_proj_weight", None) + + if bias: + self.in_proj_bias = Parameter( + torch.empty(3 * embed_dim, **factory_kwargs) + ) + else: + self.register_parameter("in_proj_bias", None) + self.out_proj = NonDynamicallyQuantizableLinear( + embed_dim, embed_dim, bias=bias, **factory_kwargs + ) + + self._reset_parameters() + else: + if not self._qkv_same_embed_dim: + raise NotImplementedError + else: + self.in_proj_linear = linear1_cls( + embed_dim, 3 * embed_dim, bias=bias, **factory_kwargs + ) + self.in_proj_weight = self.in_proj_linear.weight + + self.register_parameter("q_proj_weight", None) + self.register_parameter("k_proj_weight", None) + self.register_parameter("v_proj_weight", None) + + if bias: + self.in_proj_bias = self.in_proj_linear.bias + else: + self.register_parameter("in_proj_bias", None) + + self.out_proj = linear2_cls( + embed_dim, embed_dim, bias=bias, **factory_kwargs + ) + + if self.bias_k is not None: + xavier_normal_(self.bias_k) + if self.bias_v is not None: + xavier_normal_(self.bias_v) + + self.add_zero_attn = add_zero_attn + + def _reset_parameters(self): + if self._qkv_same_embed_dim: + xavier_uniform_(self.in_proj_weight) + else: + xavier_uniform_(self.q_proj_weight) + xavier_uniform_(self.k_proj_weight) + xavier_uniform_(self.v_proj_weight) + + if self.in_proj_bias is not None: + constant_(self.in_proj_bias, 0.0) + constant_(self.out_proj.bias, 0.0) + + if self.bias_k is not None: + xavier_normal_(self.bias_k) + if self.bias_v is not None: + xavier_normal_(self.bias_v) + + def __setstate__(self, state): + # Support loading old MultiheadAttention checkpoints generated by v1.1.0 + if "_qkv_same_embed_dim" not in state: + state["_qkv_same_embed_dim"] = True + + super(MultiheadAttention, self).__setstate__(state) + + def forward( + self, + query: Tensor, + key: Tensor, + value: Tensor, + key_padding_mask: Optional[Tensor] = None, + need_weights: bool = True, + attn_mask: Optional[Tensor] = None, + average_attn_weights: bool = True, + cache=None, + ) -> Tuple[Tensor, Optional[Tensor]]: + any_nested = query.is_nested or key.is_nested or value.is_nested + query = key = value = query.transpose(1, 0) + attn_output = multi_head_attention_forward_patched( + query, + key, + value, + self.embed_dim, + self.num_heads, + self.in_proj_weight, + self.in_proj_bias, + self.bias_k, + self.bias_v, + self.add_zero_attn, + self.dropout, + self.out_proj.weight, + self.out_proj.bias, + training=self.training, + key_padding_mask=key_padding_mask, + need_weights=need_weights, + attn_mask=attn_mask, + average_attn_weights=average_attn_weights, + cache=cache, + ) + return attn_output.transpose(1, 0) diff --git a/GPT_SoVITS/module/embedding_onnx.py b/GPT_SoVITS/module/embedding_onnx.py new file mode 100644 index 0000000..2f8dcf1 --- /dev/null +++ b/GPT_SoVITS/module/embedding_onnx.py @@ -0,0 +1,64 @@ +# modified from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/embedding.py +import math + +import torch +from torch import nn + + +class TokenEmbedding(nn.Module): + def __init__( + self, + embedding_dim: int, + vocab_size: int, + dropout: float = 0.0, + ): + super().__init__() + + self.vocab_size = vocab_size + self.embedding_dim = embedding_dim + + self.dropout = torch.nn.Dropout(p=dropout) + self.word_embeddings = nn.Embedding(self.vocab_size, self.embedding_dim) + + @property + def weight(self) -> torch.Tensor: + return self.word_embeddings.weight + + def embedding(self, index: int) -> torch.Tensor: + return self.word_embeddings.weight[index : index + 1] + + def forward(self, x: torch.Tensor): + x = self.word_embeddings(x) + x = self.dropout(x) + return x + + +class SinePositionalEmbedding(nn.Module): + def __init__( + self, + embedding_dim: int, + dropout: float = 0.0, + scale: bool = False, + alpha: bool = False, + ): + super().__init__() + self.embedding_dim = embedding_dim + self.x_scale = math.sqrt(embedding_dim) if scale else 1.0 + self.alpha = nn.Parameter(torch.ones(1), requires_grad=alpha) + self.dropout = torch.nn.Dropout(p=dropout) + self.reverse = False + self.div_term = torch.exp(torch.arange(0, self.embedding_dim, 2) * -(math.log(10000.0) / self.embedding_dim)) + self.pe = self.extend_pe(2000) + + def extend_pe(self, x): + position = torch.cumsum(torch.ones((x,1)), dim=0) + scpe = (position * self.div_term).unsqueeze(0) + pe = torch.cat([torch.sin(scpe), torch.cos(scpe)]).permute(1, 2, 0) + pe = pe.contiguous().view(1, -1, self.embedding_dim) + return pe + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pe = self.pe[:,:x.size(1),:] + output = x.unsqueeze(-1) if x.ndim == 2 else x + output = output * self.x_scale + self.alpha * pe + return self.dropout(output) diff --git a/GPT_SoVITS/module/patched_mha_with_cache_onnx.py b/GPT_SoVITS/module/patched_mha_with_cache_onnx.py new file mode 100644 index 0000000..14bdb55 --- /dev/null +++ b/GPT_SoVITS/module/patched_mha_with_cache_onnx.py @@ -0,0 +1,92 @@ +from torch.nn.functional import * +from torch.nn.functional import ( + _mha_shape_check, + _canonical_mask, + _none_or_dtype, + _in_projection_packed, +) + +def multi_head_attention_forward_patched( + query, + key, + value, + embed_dim_to_check: int, + num_heads: int, + in_proj_weight, + in_proj_bias: Optional[Tensor], + bias_k: Optional[Tensor], + bias_v: Optional[Tensor], + add_zero_attn: bool, + dropout_p: float, + out_proj_weight: Tensor, + out_proj_bias: Optional[Tensor], + training: bool = True, + key_padding_mask: Optional[Tensor] = None, + need_weights: bool = True, + attn_mask: Optional[Tensor] = None, + use_separate_proj_weight: bool = False, + q_proj_weight: Optional[Tensor] = None, + k_proj_weight: Optional[Tensor] = None, + v_proj_weight: Optional[Tensor] = None, + static_k: Optional[Tensor] = None, + static_v: Optional[Tensor] = None, + average_attn_weights: bool = True, + is_causal: bool = False, + cache=None, +) -> Tuple[Tensor, Optional[Tensor]]: + + # set up shape vars + _, _, embed_dim = query.shape + attn_mask = _canonical_mask( + mask=attn_mask, + mask_name="attn_mask", + other_type=None, + other_name="", + target_type=query.dtype, + check_other=False, + ) + head_dim = embed_dim // num_heads + + proj_qkv = linear(query, in_proj_weight, in_proj_bias) + proj_qkv = proj_qkv.unflatten(-1, (3, query.size(-1))).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous() + q, k, v = proj_qkv[0], proj_qkv[1], proj_qkv[2] + + if cache["first_infer"] == 1: + cache["k"][cache["stage"]] = k + cache["v"][cache["stage"]] = v + else: + cache["k"][cache["stage"]] = torch.cat([cache["k"][cache["stage"]][:-1], k], 0) + cache["v"][cache["stage"]] = torch.cat([cache["v"][cache["stage"]][:-1], v], 0) + k = cache["k"][cache["stage"]] + v = cache["v"][cache["stage"]] + cache["stage"] = (cache["stage"] + 1) % cache["all_stage"] + + attn_mask = _canonical_mask( + mask=attn_mask, + mask_name="attn_mask", + other_type=None, + other_name="", + target_type=q.dtype, + check_other=False, + ) + attn_mask = attn_mask.unsqueeze(0) + + q = q.view(-1, num_heads, head_dim).transpose(0, 1) + k = k.view(-1, num_heads, head_dim).transpose(0, 1) + v = v.view(-1, num_heads, head_dim).transpose(0, 1) + + dropout_p = 0.0 + attn_mask = attn_mask.unsqueeze(0) + q = q.view(num_heads, -1, head_dim).unsqueeze(0) + k = k.view(num_heads, -1, head_dim).unsqueeze(0) + v = v.view(num_heads, -1, head_dim).unsqueeze(0) + attn_output = scaled_dot_product_attention( + q, k, v, attn_mask, dropout_p, is_causal + ) + attn_output = ( + attn_output.permute(2, 0, 1, 3).contiguous().view(-1, embed_dim) + ) + attn_output = linear(attn_output, out_proj_weight, out_proj_bias) + attn_output = attn_output.view(-1, 1, attn_output.size(1)) + + return attn_output diff --git a/GPT_SoVITS/module/transformer_onnx.py b/GPT_SoVITS/module/transformer_onnx.py new file mode 100644 index 0000000..a3f68b4 --- /dev/null +++ b/GPT_SoVITS/module/transformer_onnx.py @@ -0,0 +1,292 @@ +# modified from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/transformer.py +import copy +import numbers +from functools import partial +from typing import Any +from typing import Callable +from typing import List +from typing import Optional +from typing import Tuple +from typing import Union + +import torch +from AR.modules.activation_onnx import MultiheadAttention +from AR.modules.scaling import BalancedDoubleSwish +from torch import nn +from torch import Tensor +from torch.nn import functional as F + +_shape_t = Union[int, List[int], torch.Size] + + +class LayerNorm(nn.Module): + __constants__ = ["normalized_shape", "eps", "elementwise_affine"] + normalized_shape: Tuple[int, ...] + eps: float + elementwise_affine: bool + + def __init__( + self, + normalized_shape: _shape_t, + eps: float = 1e-5, + elementwise_affine: bool = True, + device=None, + dtype=None, + ) -> None: + factory_kwargs = {"device": device, "dtype": dtype} + super(LayerNorm, self).__init__() + if isinstance(normalized_shape, numbers.Integral): + # mypy error: incompatible types in assignment + normalized_shape = (normalized_shape,) # type: ignore[assignment] + self.normalized_shape = tuple(normalized_shape) # type: ignore[arg-type] + self.eps = eps + self.elementwise_affine = elementwise_affine + if self.elementwise_affine: + self.weight = nn.Parameter( + torch.empty(self.normalized_shape, **factory_kwargs) + ) + self.bias = nn.Parameter( + torch.empty(self.normalized_shape, **factory_kwargs) + ) + else: + self.register_parameter("weight", None) + self.register_parameter("bias", None) + + self.reset_parameters() + + def reset_parameters(self) -> None: + if self.elementwise_affine: + nn.init.ones_(self.weight) + nn.init.zeros_(self.bias) + + def forward(self, input: Tensor, embedding: Any = None) -> Tensor: + if isinstance(input, tuple): + input, embedding = input + return ( + F.layer_norm( + input, + self.normalized_shape, + self.weight, + self.bias, + self.eps, + ), + embedding, + ) + + assert embedding is None + return F.layer_norm( + input, self.normalized_shape, self.weight, self.bias, self.eps + ) + + def extra_repr(self) -> str: + return ( + "{normalized_shape}, eps={eps}, " + "elementwise_affine={elementwise_affine}".format(**self.__dict__) + ) + + +class IdentityNorm(nn.Module): + def __init__( + self, + d_model: int, + eps: float = 1e-5, + device=None, + dtype=None, + ) -> None: + super(IdentityNorm, self).__init__() + + def forward(self, input: Tensor, embedding: Any = None) -> Tensor: + if isinstance(input, tuple): + return input + + assert embedding is None + return input + + +class TransformerEncoder(nn.Module): + r"""TransformerEncoder is a stack of N encoder layers. Users can build the + BERT(https://arxiv.org/abs/1810.04805) model with corresponding parameters. + + Args: + encoder_layer: an instance of the TransformerEncoderLayer() class (required). + num_layers: the number of sub-encoder-layers in the encoder (required). + norm: the layer normalization component (optional). + enable_nested_tensor: if True, input will automatically convert to nested tensor + (and convert back on output). This will improve the overall performance of + TransformerEncoder when padding rate is high. Default: ``True`` (enabled). + + Examples:: + >>> encoder_layer = TransformerEncoderLayer(d_model=512, nhead=8) + >>> transformer_encoder = TransformerEncoder(encoder_layer, num_layers=6) + >>> src = torch.rand(10, 32, 512) + >>> out = transformer_encoder(src) + """ + __constants__ = ["norm"] + + def __init__(self, encoder_layer, num_layers, norm=None): + super(TransformerEncoder, self).__init__() + self.layers = _get_clones(encoder_layer, num_layers) + self.num_layers = num_layers + self.norm = norm + + def forward( + self, + src: Tensor, + mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + return_layer_states: bool = False, + cache=None, + ) -> Tensor: + output = src + for mod in self.layers: + output = mod( + output, + src_mask=mask, + src_key_padding_mask=src_key_padding_mask, + cache=cache, + ) + + if self.norm is not None: + output = self.norm(output) + + return output + + +class TransformerEncoderLayer(nn.Module): + __constants__ = ["batch_first", "norm_first"] + def __init__( + self, + d_model: int, + nhead: int, + dim_feedforward: int = 2048, + dropout: float = 0.1, + activation: Union[str, Callable[[Tensor], Tensor]] = F.relu, + batch_first: bool = False, + norm_first: bool = False, + device=None, + dtype=None, + linear1_self_attention_cls: nn.Module = nn.Linear, + linear2_self_attention_cls: nn.Module = nn.Linear, + linear1_feedforward_cls: nn.Module = nn.Linear, + linear2_feedforward_cls: nn.Module = nn.Linear, + layer_norm_cls: nn.Module = LayerNorm, + layer_norm_eps: float = 1e-5, + adaptive_layer_norm=False, + ) -> None: + factory_kwargs = {"device": device, "dtype": dtype} + super(TransformerEncoderLayer, self).__init__() + self.self_attn = MultiheadAttention( + d_model, # 512 16 + nhead, + dropout=dropout, + batch_first=batch_first, + linear1_cls=linear1_self_attention_cls, + linear2_cls=linear2_self_attention_cls, + **factory_kwargs, + ) + self.linear1 = linear1_feedforward_cls( + d_model, dim_feedforward, **factory_kwargs + ) + self.dropout = nn.Dropout(dropout) + self.linear2 = linear2_feedforward_cls( + dim_feedforward, d_model, **factory_kwargs + ) + self.norm_first = norm_first + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + if isinstance(activation, str): + activation = _get_activation_fn(activation) + elif isinstance(activation, partial): + activation = activation(d_model) + elif activation == BalancedDoubleSwish: + activation = BalancedDoubleSwish(d_model) + self.activation = activation + + norm1 = layer_norm_cls(d_model, eps=layer_norm_eps, **factory_kwargs) + if layer_norm_cls == IdentityNorm: + norm2 = BalancedBasicNorm(d_model, eps=layer_norm_eps, **factory_kwargs) + else: + norm2 = layer_norm_cls(d_model, eps=layer_norm_eps, **factory_kwargs) + + if adaptive_layer_norm: + self.norm1 = AdaptiveLayerNorm(d_model, norm1) + self.norm2 = AdaptiveLayerNorm(d_model, norm2) + else: + self.norm1 = norm1 + self.norm2 = norm2 + + def __setstate__(self, state): + super(TransformerEncoderLayer, self).__setstate__(state) + if not hasattr(self, "activation"): + self.activation = F.relu + + def forward( + self, + src: Tensor, + src_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + cache=None, + ) -> Tensor: + x = src + stage_embedding = None + x = self.norm1( + x + self._sa_block(x, src_mask, src_key_padding_mask, cache=cache), + stage_embedding, + ) + x = self.norm2(x + self._ff_block(x), stage_embedding) + + return x + + def _sa_block( + self, + x: Tensor, + attn_mask: Optional[Tensor], + key_padding_mask: Optional[Tensor], + cache=None, + ) -> Tensor: + x = self.self_attn( + x, + x, + x, + attn_mask=attn_mask, + key_padding_mask=key_padding_mask, + need_weights=False, + cache=cache, + ) + return self.dropout1(x) + + def _ff_block(self, x: Tensor) -> Tensor: + x = self.linear2(self.dropout(self.activation(self.linear1(x)))) + return self.dropout2(x) + + +class AdaptiveLayerNorm(nn.Module): + r"""Adaptive Layer Normalization""" + + def __init__(self, d_model, norm) -> None: + super(AdaptiveLayerNorm, self).__init__() + self.project_layer = nn.Linear(d_model, 2 * d_model) + self.norm = norm + self.d_model = d_model + self.eps = self.norm.eps + + def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor: + if isinstance(input, tuple): + input, embedding = input + weight, bias = torch.split( + self.project_layer(embedding), + split_size_or_sections=self.d_model, + dim=-1, + ) + return (weight * self.norm(input) + bias, embedding) + + weight, bias = torch.split( + self.project_layer(embedding), + split_size_or_sections=self.d_model, + dim=-1, + ) + return weight * self.norm(input) + bias + + +def _get_clones(module, N): + return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) From aad68f716d67e4821fdad14736c6a92914e5af98 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CE=9D=CE=B1=CF=81=CE=BF=CF=85=CF=83=CE=AD=C2=B7=CE=BC?= =?UTF-8?q?=C2=B7=CE=B3=CE=B9=CE=BF=CF=85=CE=BC=CE=B5=CE=BC=CE=AF=C2=B7?= =?UTF-8?q?=CE=A7=CE=B9=CE=BD=CE=B1=CE=BA=CE=AC=CE=BD=CE=BD=CE=B1?= <40709280+NaruseMioShirakana@users.noreply.github.com> Date: Sat, 6 Apr 2024 15:47:35 +0800 Subject: [PATCH 02/14] Delete GPT_SoVITS/module/activation_onnx.py --- GPT_SoVITS/module/activation_onnx.py | 178 --------------------------- 1 file changed, 178 deletions(-) delete mode 100644 GPT_SoVITS/module/activation_onnx.py diff --git a/GPT_SoVITS/module/activation_onnx.py b/GPT_SoVITS/module/activation_onnx.py deleted file mode 100644 index b54acd9..0000000 --- a/GPT_SoVITS/module/activation_onnx.py +++ /dev/null @@ -1,178 +0,0 @@ -# modified from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/activation.py -from typing import Optional -from typing import Tuple -import torch -from torch import Tensor -from torch.nn import Linear -from torch.nn import Module -from torch.nn.init import constant_ -from torch.nn.init import xavier_normal_ -from torch.nn.init import xavier_uniform_ -from torch.nn.modules.linear import NonDynamicallyQuantizableLinear -from torch.nn.parameter import Parameter - -from torch.nn import functional as F -from AR.modules.patched_mha_with_cache_onnx import multi_head_attention_forward_patched - - -class MultiheadAttention(Module): - __constants__ = ["batch_first"] - bias_k: Optional[torch.Tensor] - bias_v: Optional[torch.Tensor] - - def __init__( - self, - embed_dim, - num_heads, - dropout=0.0, - bias=True, - add_bias_kv=False, - add_zero_attn=False, - kdim=None, - vdim=None, - batch_first=False, - linear1_cls=Linear, - linear2_cls=Linear, - device=None, - dtype=None, - ) -> None: - factory_kwargs = {"device": device, "dtype": dtype} - super(MultiheadAttention, self).__init__() - self.embed_dim = embed_dim - self.kdim = kdim if kdim is not None else embed_dim - self.vdim = vdim if vdim is not None else embed_dim - self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim - - self.num_heads = num_heads - self.dropout = dropout - self.batch_first = batch_first - self.head_dim = embed_dim // num_heads - assert ( - self.head_dim * num_heads == self.embed_dim - ), "embed_dim must be divisible by num_heads" - - if add_bias_kv: - self.bias_k = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs)) - self.bias_v = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs)) - else: - self.bias_k = self.bias_v = None - - if linear1_cls == Linear: - if not self._qkv_same_embed_dim: - self.q_proj_weight = Parameter( - torch.empty((embed_dim, embed_dim), **factory_kwargs) - ) - self.k_proj_weight = Parameter( - torch.empty((embed_dim, self.kdim), **factory_kwargs) - ) - self.v_proj_weight = Parameter( - torch.empty((embed_dim, self.vdim), **factory_kwargs) - ) - self.register_parameter("in_proj_weight", None) - else: - self.in_proj_weight = Parameter( - torch.empty((3 * embed_dim, embed_dim), **factory_kwargs) - ) - self.register_parameter("q_proj_weight", None) - self.register_parameter("k_proj_weight", None) - self.register_parameter("v_proj_weight", None) - - if bias: - self.in_proj_bias = Parameter( - torch.empty(3 * embed_dim, **factory_kwargs) - ) - else: - self.register_parameter("in_proj_bias", None) - self.out_proj = NonDynamicallyQuantizableLinear( - embed_dim, embed_dim, bias=bias, **factory_kwargs - ) - - self._reset_parameters() - else: - if not self._qkv_same_embed_dim: - raise NotImplementedError - else: - self.in_proj_linear = linear1_cls( - embed_dim, 3 * embed_dim, bias=bias, **factory_kwargs - ) - self.in_proj_weight = self.in_proj_linear.weight - - self.register_parameter("q_proj_weight", None) - self.register_parameter("k_proj_weight", None) - self.register_parameter("v_proj_weight", None) - - if bias: - self.in_proj_bias = self.in_proj_linear.bias - else: - self.register_parameter("in_proj_bias", None) - - self.out_proj = linear2_cls( - embed_dim, embed_dim, bias=bias, **factory_kwargs - ) - - if self.bias_k is not None: - xavier_normal_(self.bias_k) - if self.bias_v is not None: - xavier_normal_(self.bias_v) - - self.add_zero_attn = add_zero_attn - - def _reset_parameters(self): - if self._qkv_same_embed_dim: - xavier_uniform_(self.in_proj_weight) - else: - xavier_uniform_(self.q_proj_weight) - xavier_uniform_(self.k_proj_weight) - xavier_uniform_(self.v_proj_weight) - - if self.in_proj_bias is not None: - constant_(self.in_proj_bias, 0.0) - constant_(self.out_proj.bias, 0.0) - - if self.bias_k is not None: - xavier_normal_(self.bias_k) - if self.bias_v is not None: - xavier_normal_(self.bias_v) - - def __setstate__(self, state): - # Support loading old MultiheadAttention checkpoints generated by v1.1.0 - if "_qkv_same_embed_dim" not in state: - state["_qkv_same_embed_dim"] = True - - super(MultiheadAttention, self).__setstate__(state) - - def forward( - self, - query: Tensor, - key: Tensor, - value: Tensor, - key_padding_mask: Optional[Tensor] = None, - need_weights: bool = True, - attn_mask: Optional[Tensor] = None, - average_attn_weights: bool = True, - cache=None, - ) -> Tuple[Tensor, Optional[Tensor]]: - any_nested = query.is_nested or key.is_nested or value.is_nested - query = key = value = query.transpose(1, 0) - attn_output = multi_head_attention_forward_patched( - query, - key, - value, - self.embed_dim, - self.num_heads, - self.in_proj_weight, - self.in_proj_bias, - self.bias_k, - self.bias_v, - self.add_zero_attn, - self.dropout, - self.out_proj.weight, - self.out_proj.bias, - training=self.training, - key_padding_mask=key_padding_mask, - need_weights=need_weights, - attn_mask=attn_mask, - average_attn_weights=average_attn_weights, - cache=cache, - ) - return attn_output.transpose(1, 0) From 405f115dffe97bfd0b9a5ef76acc172bdccc5c73 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CE=9D=CE=B1=CF=81=CE=BF=CF=85=CF=83=CE=AD=C2=B7=CE=BC?= =?UTF-8?q?=C2=B7=CE=B3=CE=B9=CE=BF=CF=85=CE=BC=CE=B5=CE=BC=CE=AF=C2=B7?= =?UTF-8?q?=CE=A7=CE=B9=CE=BD=CE=B1=CE=BA=CE=AC=CE=BD=CE=BD=CE=B1?= <40709280+NaruseMioShirakana@users.noreply.github.com> Date: Sat, 6 Apr 2024 15:47:44 +0800 Subject: [PATCH 03/14] Delete GPT_SoVITS/module/attentions_onnx.py --- GPT_SoVITS/module/attentions_onnx.py | 354 --------------------------- 1 file changed, 354 deletions(-) delete mode 100644 GPT_SoVITS/module/attentions_onnx.py diff --git a/GPT_SoVITS/module/attentions_onnx.py b/GPT_SoVITS/module/attentions_onnx.py deleted file mode 100644 index bc63a06..0000000 --- a/GPT_SoVITS/module/attentions_onnx.py +++ /dev/null @@ -1,354 +0,0 @@ -import math -import torch -from torch import nn -from torch.nn import functional as F - -from module import commons -from module.modules import LayerNorm - - -class LayerNorm(nn.Module): - def __init__(self, channels, eps=1e-5): - super().__init__() - self.channels = channels - self.eps = eps - - self.gamma = nn.Parameter(torch.ones(channels)) - self.beta = nn.Parameter(torch.zeros(channels)) - - def forward(self, x): - x = x.transpose(1, -1) - x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps) - return x.transpose(1, -1) - - -@torch.jit.script -def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels): - n_channels_int = n_channels[0] - in_act = input_a + input_b - t_act = torch.tanh(in_act[:, :n_channels_int, :]) - s_act = torch.sigmoid(in_act[:, n_channels_int:, :]) - acts = t_act * s_act - return acts - - -class Encoder(nn.Module): - def __init__( - self, - hidden_channels, - filter_channels, - n_heads, - n_layers, - kernel_size=1, - p_dropout=0.0, - window_size=4, - isflow=True, - **kwargs - ): - super().__init__() - self.hidden_channels = hidden_channels - self.filter_channels = filter_channels - self.n_heads = n_heads - self.n_layers = n_layers - self.kernel_size = kernel_size - self.p_dropout = p_dropout - self.window_size = window_size - # if isflow: - # cond_layer = torch.nn.Conv1d(256, 2*hidden_channels*n_layers, 1) - # self.cond_pre = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, 1) - # self.cond_layer = weight_norm(cond_layer, name='weight') - # self.gin_channels = 256 - self.cond_layer_idx = self.n_layers - if "gin_channels" in kwargs: - self.gin_channels = kwargs["gin_channels"] - if self.gin_channels != 0: - self.spk_emb_linear = nn.Linear(self.gin_channels, self.hidden_channels) - # vits2 says 3rd block, so idx is 2 by default - self.cond_layer_idx = ( - kwargs["cond_layer_idx"] if "cond_layer_idx" in kwargs else 2 - ) - logging.debug(self.gin_channels, self.cond_layer_idx) - assert ( - self.cond_layer_idx < self.n_layers - ), "cond_layer_idx should be less than n_layers" - self.drop = nn.Dropout(p_dropout) - self.attn_layers = nn.ModuleList() - self.norm_layers_1 = nn.ModuleList() - self.ffn_layers = nn.ModuleList() - self.norm_layers_2 = nn.ModuleList() - for i in range(self.n_layers): - self.attn_layers.append( - MultiHeadAttention( - hidden_channels, - hidden_channels, - n_heads, - p_dropout=p_dropout, - window_size=window_size, - ) - ) - self.norm_layers_1.append(LayerNorm(hidden_channels)) - self.ffn_layers.append( - FFN( - hidden_channels, - hidden_channels, - filter_channels, - kernel_size, - p_dropout=p_dropout, - ) - ) - self.norm_layers_2.append(LayerNorm(hidden_channels)) - - def forward(self, x, x_mask, g=None): - attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1) - x = x * x_mask - for i in range(self.n_layers): - if i == self.cond_layer_idx and g is not None: - g = self.spk_emb_linear(g.transpose(1, 2)) - g = g.transpose(1, 2) - x = x + g - x = x * x_mask - y = self.attn_layers[i](x, x, attn_mask) - y = self.drop(y) - x = self.norm_layers_1[i](x + y) - - y = self.ffn_layers[i](x, x_mask) - y = self.drop(y) - x = self.norm_layers_2[i](x + y) - x = x * x_mask - return x - - -class MultiHeadAttention(nn.Module): - def __init__( - self, - channels, - out_channels, - n_heads, - p_dropout=0.0, - window_size=None, - heads_share=True, - block_length=None, - proximal_bias=False, - proximal_init=False, - ): - super().__init__() - assert channels % n_heads == 0 - - self.channels = channels - self.out_channels = out_channels - self.n_heads = n_heads - self.p_dropout = p_dropout - self.window_size = window_size - self.heads_share = heads_share - self.block_length = block_length - self.proximal_bias = proximal_bias - self.proximal_init = proximal_init - self.attn = None - - self.k_channels = channels // n_heads - self.conv_q = nn.Conv1d(channels, channels, 1) - self.conv_k = nn.Conv1d(channels, channels, 1) - self.conv_v = nn.Conv1d(channels, channels, 1) - self.conv_o = nn.Conv1d(channels, out_channels, 1) - self.drop = nn.Dropout(p_dropout) - - if window_size is not None: - n_heads_rel = 1 if heads_share else n_heads - rel_stddev = self.k_channels**-0.5 - self.emb_rel_k = nn.Parameter( - torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) - * rel_stddev - ) - self.emb_rel_v = nn.Parameter( - torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) - * rel_stddev - ) - - nn.init.xavier_uniform_(self.conv_q.weight) - nn.init.xavier_uniform_(self.conv_k.weight) - nn.init.xavier_uniform_(self.conv_v.weight) - if proximal_init: - with torch.no_grad(): - self.conv_k.weight.copy_(self.conv_q.weight) - self.conv_k.bias.copy_(self.conv_q.bias) - - def forward(self, x, c, attn_mask=None): - q = self.conv_q(x) - k = self.conv_k(c) - v = self.conv_v(c) - - x, self.attn = self.attention(q, k, v, mask=attn_mask) - - x = self.conv_o(x) - return x - - def attention(self, query, key, value, mask=None): - # reshape [b, d, t] -> [b, n_h, t, d_k] - b, d, t_s, _ = (*key.size(), query.size(2)) - query = query.view(b, self.n_heads, self.k_channels, -1).transpose(2, 3) - key = key.view(b, self.n_heads, self.k_channels, -1).transpose(2, 3) - value = value.view(b, self.n_heads, self.k_channels, -1).transpose(2, 3) - scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1)) - - if self.window_size is not None: - key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s) - rel_logits = self._matmul_with_relative_keys(query / math.sqrt(self.k_channels), key_relative_embeddings) - scores_local = self._relative_position_to_absolute_position(rel_logits) - scores = scores + scores_local - - if mask is not None: - scores = scores.masked_fill(mask == 0, -1e4) - - p_attn = F.softmax(scores, dim=-1) - p_attn = self.drop(p_attn) - output = torch.matmul(p_attn, value) - - if self.window_size is not None: - relative_weights = self._absolute_position_to_relative_position(p_attn) - value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, t_s) - output = output + self._matmul_with_relative_values(relative_weights, value_relative_embeddings) - - output = (output.transpose(2, 3).contiguous().view(b, d, -1)) - return output, p_attn - - def _matmul_with_relative_values(self, x, y): - """ - x: [b, h, l, m] - y: [h or 1, m, d] - ret: [b, h, l, d] - """ - ret = torch.matmul(x, y.unsqueeze(0)) - return ret - - def _matmul_with_relative_keys(self, x, y): - """ - x: [b, h, l, d] - y: [h or 1, m, d] - ret: [b, h, l, m] - """ - ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1)) - return ret - - def _get_relative_embeddings(self, relative_embeddings, length): - max_relative_position = 2 * self.window_size + 1 - # Pad first before slice to avoid using cond ops. - pad_l = torch.zeros((1), dtype = torch.int64) + length - (self.window_size + 1) - pad_s = torch.zeros((1), dtype = torch.int64) + (self.window_size + 1) - length - pad_length = torch.max(pad_l, other=torch.zeros((1), dtype = torch.int64)) - slice_start_position = torch.max(pad_s, other=torch.zeros((1), dtype = torch.int64)) - - slice_end_position = slice_start_position + 2 * length - 1 - padded_relative_embeddings = F.pad( - relative_embeddings, - commons.convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]), - ) - used_relative_embeddings = padded_relative_embeddings[ - :, slice_start_position:slice_end_position - ] - return used_relative_embeddings - - def _relative_position_to_absolute_position(self, x): - """ - x: [b, h, l, 2*l-1] - ret: [b, h, l, l] - """ - batch, heads, length, _ = x.size() - # Concat columns of pad to shift from relative to absolute indexing. - x = F.pad(x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, 1]])) - - # Concat extra elements so to add up to shape (len+1, 2*len-1). - x_flat = x.view([batch, heads, length * 2 * length]) - x_flat = F.pad( - x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [0, length - 1]]) - ) - - # Reshape and slice out the padded elements. - x_final = x_flat.view([batch, heads, length + 1, 2 * length - 1])[ - :, :, :length, length - 1 : - ] - return x_final - - def _absolute_position_to_relative_position(self, x): - """ - x: [b, h, l, l] - ret: [b, h, l, 2*l-1] - """ - batch, heads, length, _ = x.size() - # padd along column - x = F.pad( - x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length - 1]]) - ) - x_flat = x.view([batch, heads, length**2 + length * (length - 1)]) - # add 0's in the beginning that will skew the elements after reshape - x_flat = F.pad(x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [length, 0]])) - x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:] - return x_final - - def _attention_bias_proximal(self, length): - """Bias for self-attention to encourage attention to close positions. - Args: - length: an integer scalar. - Returns: - a Tensor with shape [1, 1, length, length] - """ - r = torch.arange(length, dtype=torch.float32) - diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1) - return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0) - - -class FFN(nn.Module): - def __init__( - self, - in_channels, - out_channels, - filter_channels, - kernel_size, - p_dropout=0.0, - activation=None, - causal=False, - ): - super().__init__() - self.in_channels = in_channels - self.out_channels = out_channels - self.filter_channels = filter_channels - self.kernel_size = kernel_size - self.p_dropout = p_dropout - self.activation = activation - self.causal = causal - - if causal: - self.padding = self._causal_padding - else: - self.padding = self._same_padding - - self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size) - self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size) - self.drop = nn.Dropout(p_dropout) - - def forward(self, x, x_mask): - x = self.conv_1(self.padding(x * x_mask)) - if self.activation == "gelu": - x = x * torch.sigmoid(1.702 * x) - else: - x = torch.relu(x) - x = self.drop(x) - x = self.conv_2(self.padding(x * x_mask)) - return x * x_mask - - def _causal_padding(self, x): - if self.kernel_size == 1: - return x - pad_l = self.kernel_size - 1 - pad_r = 0 - padding = [[0, 0], [0, 0], [pad_l, pad_r]] - x = F.pad(x, commons.convert_pad_shape(padding)) - return x - - def _same_padding(self, x): - if self.kernel_size == 1: - return x - pad_l = (self.kernel_size - 1) // 2 - pad_r = self.kernel_size // 2 - padding = [[0, 0], [0, 0], [pad_l, pad_r]] - x = F.pad(x, commons.convert_pad_shape(padding)) - return x From 96ff0008a4775e9f08d2985a92cd374ca4e001e1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CE=9D=CE=B1=CF=81=CE=BF=CF=85=CF=83=CE=AD=C2=B7=CE=BC?= =?UTF-8?q?=C2=B7=CE=B3=CE=B9=CE=BF=CF=85=CE=BC=CE=B5=CE=BC=CE=AF=C2=B7?= =?UTF-8?q?=CE=A7=CE=B9=CE=BD=CE=B1=CE=BA=CE=AC=CE=BD=CE=BD=CE=B1?= <40709280+NaruseMioShirakana@users.noreply.github.com> Date: Sat, 6 Apr 2024 15:47:53 +0800 Subject: [PATCH 04/14] Delete GPT_SoVITS/module/embedding_onnx.py --- GPT_SoVITS/module/embedding_onnx.py | 64 ----------------------------- 1 file changed, 64 deletions(-) delete mode 100644 GPT_SoVITS/module/embedding_onnx.py diff --git a/GPT_SoVITS/module/embedding_onnx.py b/GPT_SoVITS/module/embedding_onnx.py deleted file mode 100644 index 2f8dcf1..0000000 --- a/GPT_SoVITS/module/embedding_onnx.py +++ /dev/null @@ -1,64 +0,0 @@ -# modified from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/embedding.py -import math - -import torch -from torch import nn - - -class TokenEmbedding(nn.Module): - def __init__( - self, - embedding_dim: int, - vocab_size: int, - dropout: float = 0.0, - ): - super().__init__() - - self.vocab_size = vocab_size - self.embedding_dim = embedding_dim - - self.dropout = torch.nn.Dropout(p=dropout) - self.word_embeddings = nn.Embedding(self.vocab_size, self.embedding_dim) - - @property - def weight(self) -> torch.Tensor: - return self.word_embeddings.weight - - def embedding(self, index: int) -> torch.Tensor: - return self.word_embeddings.weight[index : index + 1] - - def forward(self, x: torch.Tensor): - x = self.word_embeddings(x) - x = self.dropout(x) - return x - - -class SinePositionalEmbedding(nn.Module): - def __init__( - self, - embedding_dim: int, - dropout: float = 0.0, - scale: bool = False, - alpha: bool = False, - ): - super().__init__() - self.embedding_dim = embedding_dim - self.x_scale = math.sqrt(embedding_dim) if scale else 1.0 - self.alpha = nn.Parameter(torch.ones(1), requires_grad=alpha) - self.dropout = torch.nn.Dropout(p=dropout) - self.reverse = False - self.div_term = torch.exp(torch.arange(0, self.embedding_dim, 2) * -(math.log(10000.0) / self.embedding_dim)) - self.pe = self.extend_pe(2000) - - def extend_pe(self, x): - position = torch.cumsum(torch.ones((x,1)), dim=0) - scpe = (position * self.div_term).unsqueeze(0) - pe = torch.cat([torch.sin(scpe), torch.cos(scpe)]).permute(1, 2, 0) - pe = pe.contiguous().view(1, -1, self.embedding_dim) - return pe - - def forward(self, x: torch.Tensor) -> torch.Tensor: - pe = self.pe[:,:x.size(1),:] - output = x.unsqueeze(-1) if x.ndim == 2 else x - output = output * self.x_scale + self.alpha * pe - return self.dropout(output) From a17399d75b8c5efb6aa9cb8a0c8d3c8590df080b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CE=9D=CE=B1=CF=81=CE=BF=CF=85=CF=83=CE=AD=C2=B7=CE=BC?= =?UTF-8?q?=C2=B7=CE=B3=CE=B9=CE=BF=CF=85=CE=BC=CE=B5=CE=BC=CE=AF=C2=B7?= =?UTF-8?q?=CE=A7=CE=B9=CE=BD=CE=B1=CE=BA=CE=AC=CE=BD=CE=BD=CE=B1?= <40709280+NaruseMioShirakana@users.noreply.github.com> Date: Sat, 6 Apr 2024 15:48:03 +0800 Subject: [PATCH 05/14] Delete GPT_SoVITS/module/patched_mha_with_cache_onnx.py --- .../module/patched_mha_with_cache_onnx.py | 92 ------------------- 1 file changed, 92 deletions(-) delete mode 100644 GPT_SoVITS/module/patched_mha_with_cache_onnx.py diff --git a/GPT_SoVITS/module/patched_mha_with_cache_onnx.py b/GPT_SoVITS/module/patched_mha_with_cache_onnx.py deleted file mode 100644 index 14bdb55..0000000 --- a/GPT_SoVITS/module/patched_mha_with_cache_onnx.py +++ /dev/null @@ -1,92 +0,0 @@ -from torch.nn.functional import * -from torch.nn.functional import ( - _mha_shape_check, - _canonical_mask, - _none_or_dtype, - _in_projection_packed, -) - -def multi_head_attention_forward_patched( - query, - key, - value, - embed_dim_to_check: int, - num_heads: int, - in_proj_weight, - in_proj_bias: Optional[Tensor], - bias_k: Optional[Tensor], - bias_v: Optional[Tensor], - add_zero_attn: bool, - dropout_p: float, - out_proj_weight: Tensor, - out_proj_bias: Optional[Tensor], - training: bool = True, - key_padding_mask: Optional[Tensor] = None, - need_weights: bool = True, - attn_mask: Optional[Tensor] = None, - use_separate_proj_weight: bool = False, - q_proj_weight: Optional[Tensor] = None, - k_proj_weight: Optional[Tensor] = None, - v_proj_weight: Optional[Tensor] = None, - static_k: Optional[Tensor] = None, - static_v: Optional[Tensor] = None, - average_attn_weights: bool = True, - is_causal: bool = False, - cache=None, -) -> Tuple[Tensor, Optional[Tensor]]: - - # set up shape vars - _, _, embed_dim = query.shape - attn_mask = _canonical_mask( - mask=attn_mask, - mask_name="attn_mask", - other_type=None, - other_name="", - target_type=query.dtype, - check_other=False, - ) - head_dim = embed_dim // num_heads - - proj_qkv = linear(query, in_proj_weight, in_proj_bias) - proj_qkv = proj_qkv.unflatten(-1, (3, query.size(-1))).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous() - q, k, v = proj_qkv[0], proj_qkv[1], proj_qkv[2] - - if cache["first_infer"] == 1: - cache["k"][cache["stage"]] = k - cache["v"][cache["stage"]] = v - else: - cache["k"][cache["stage"]] = torch.cat([cache["k"][cache["stage"]][:-1], k], 0) - cache["v"][cache["stage"]] = torch.cat([cache["v"][cache["stage"]][:-1], v], 0) - k = cache["k"][cache["stage"]] - v = cache["v"][cache["stage"]] - cache["stage"] = (cache["stage"] + 1) % cache["all_stage"] - - attn_mask = _canonical_mask( - mask=attn_mask, - mask_name="attn_mask", - other_type=None, - other_name="", - target_type=q.dtype, - check_other=False, - ) - attn_mask = attn_mask.unsqueeze(0) - - q = q.view(-1, num_heads, head_dim).transpose(0, 1) - k = k.view(-1, num_heads, head_dim).transpose(0, 1) - v = v.view(-1, num_heads, head_dim).transpose(0, 1) - - dropout_p = 0.0 - attn_mask = attn_mask.unsqueeze(0) - q = q.view(num_heads, -1, head_dim).unsqueeze(0) - k = k.view(num_heads, -1, head_dim).unsqueeze(0) - v = v.view(num_heads, -1, head_dim).unsqueeze(0) - attn_output = scaled_dot_product_attention( - q, k, v, attn_mask, dropout_p, is_causal - ) - attn_output = ( - attn_output.permute(2, 0, 1, 3).contiguous().view(-1, embed_dim) - ) - attn_output = linear(attn_output, out_proj_weight, out_proj_bias) - attn_output = attn_output.view(-1, 1, attn_output.size(1)) - - return attn_output From 7c3492c2406f51e9048a57985f590b78e888abde Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CE=9D=CE=B1=CF=81=CE=BF=CF=85=CF=83=CE=AD=C2=B7=CE=BC?= =?UTF-8?q?=C2=B7=CE=B3=CE=B9=CE=BF=CF=85=CE=BC=CE=B5=CE=BC=CE=AF=C2=B7?= =?UTF-8?q?=CE=A7=CE=B9=CE=BD=CE=B1=CE=BA=CE=AC=CE=BD=CE=BD=CE=B1?= <40709280+NaruseMioShirakana@users.noreply.github.com> Date: Sat, 6 Apr 2024 15:48:36 +0800 Subject: [PATCH 06/14] Add files via upload --- GPT_SoVITS/AR/modules/embedding_onnx.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/GPT_SoVITS/AR/modules/embedding_onnx.py b/GPT_SoVITS/AR/modules/embedding_onnx.py index b93405b..2f8dcf1 100644 --- a/GPT_SoVITS/AR/modules/embedding_onnx.py +++ b/GPT_SoVITS/AR/modules/embedding_onnx.py @@ -48,16 +48,17 @@ class SinePositionalEmbedding(nn.Module): self.dropout = torch.nn.Dropout(p=dropout) self.reverse = False self.div_term = torch.exp(torch.arange(0, self.embedding_dim, 2) * -(math.log(10000.0) / self.embedding_dim)) + self.pe = self.extend_pe(2000) def extend_pe(self, x): - position = torch.cumsum(torch.ones_like(x[:,:,0]), dim=1).transpose(0, 1) + position = torch.cumsum(torch.ones((x,1)), dim=0) scpe = (position * self.div_term).unsqueeze(0) pe = torch.cat([torch.sin(scpe), torch.cos(scpe)]).permute(1, 2, 0) pe = pe.contiguous().view(1, -1, self.embedding_dim) return pe def forward(self, x: torch.Tensor) -> torch.Tensor: - pe = self.extend_pe(x) + pe = self.pe[:,:x.size(1),:] output = x.unsqueeze(-1) if x.ndim == 2 else x output = output * self.x_scale + self.alpha * pe return self.dropout(output) From dd8bd6000ca00bfe5875e6f8c247ac2e403c0235 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CE=9D=CE=B1=CF=81=CE=BF=CF=85=CF=83=CE=AD=C2=B7=CE=BC?= =?UTF-8?q?=C2=B7=CE=B3=CE=B9=CE=BF=CF=85=CE=BC=CE=B5=CE=BC=CE=AF=C2=B7?= =?UTF-8?q?=CE=A7=CE=B9=CE=BD=CE=B1=CE=BA=CE=AC=CE=BD=CE=BD=CE=B1?= <40709280+NaruseMioShirakana@users.noreply.github.com> Date: Sat, 6 Apr 2024 15:49:02 +0800 Subject: [PATCH 07/14] Add files via upload --- GPT_SoVITS/AR/models/t2s_lightning_module_onnx.py | 3 +-- GPT_SoVITS/AR/models/t2s_model_onnx.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/GPT_SoVITS/AR/models/t2s_lightning_module_onnx.py b/GPT_SoVITS/AR/models/t2s_lightning_module_onnx.py index 487edb0..bb9e30b 100644 --- a/GPT_SoVITS/AR/models/t2s_lightning_module_onnx.py +++ b/GPT_SoVITS/AR/models/t2s_lightning_module_onnx.py @@ -1,5 +1,4 @@ -# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/models/t2s_lightning_module.py -# reference: https://github.com/lifeiteng/vall-e +# modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/model/t2s_lightning_module.py import os, sys now_dir = os.getcwd() diff --git a/GPT_SoVITS/AR/models/t2s_model_onnx.py b/GPT_SoVITS/AR/models/t2s_model_onnx.py index 7834297..92f2d74 100644 --- a/GPT_SoVITS/AR/models/t2s_model_onnx.py +++ b/GPT_SoVITS/AR/models/t2s_model_onnx.py @@ -1,5 +1,4 @@ -# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/models/t2s_model.py -# reference: https://github.com/lifeiteng/vall-e +# modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/model/t2s_model.py import torch from tqdm import tqdm From d89ff2d2c9f87dd953decbd69638f3f530690128 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CE=9D=CE=B1=CF=81=CE=BF=CF=85=CF=83=CE=AD=C2=B7=CE=BC?= =?UTF-8?q?=C2=B7=CE=B3=CE=B9=CE=BF=CF=85=CE=BC=CE=B5=CE=BC=CE=AF=C2=B7?= =?UTF-8?q?=CE=A7=CE=B9=CE=BD=CE=B1=CE=BA=CE=AC=CE=BD=CE=BD=CE=B1?= <40709280+NaruseMioShirakana@users.noreply.github.com> Date: Sat, 6 Apr 2024 15:50:07 +0800 Subject: [PATCH 08/14] Delete GPT_SoVITS/module/transformer_onnx.py --- GPT_SoVITS/module/transformer_onnx.py | 292 -------------------------- 1 file changed, 292 deletions(-) delete mode 100644 GPT_SoVITS/module/transformer_onnx.py diff --git a/GPT_SoVITS/module/transformer_onnx.py b/GPT_SoVITS/module/transformer_onnx.py deleted file mode 100644 index a3f68b4..0000000 --- a/GPT_SoVITS/module/transformer_onnx.py +++ /dev/null @@ -1,292 +0,0 @@ -# modified from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/transformer.py -import copy -import numbers -from functools import partial -from typing import Any -from typing import Callable -from typing import List -from typing import Optional -from typing import Tuple -from typing import Union - -import torch -from AR.modules.activation_onnx import MultiheadAttention -from AR.modules.scaling import BalancedDoubleSwish -from torch import nn -from torch import Tensor -from torch.nn import functional as F - -_shape_t = Union[int, List[int], torch.Size] - - -class LayerNorm(nn.Module): - __constants__ = ["normalized_shape", "eps", "elementwise_affine"] - normalized_shape: Tuple[int, ...] - eps: float - elementwise_affine: bool - - def __init__( - self, - normalized_shape: _shape_t, - eps: float = 1e-5, - elementwise_affine: bool = True, - device=None, - dtype=None, - ) -> None: - factory_kwargs = {"device": device, "dtype": dtype} - super(LayerNorm, self).__init__() - if isinstance(normalized_shape, numbers.Integral): - # mypy error: incompatible types in assignment - normalized_shape = (normalized_shape,) # type: ignore[assignment] - self.normalized_shape = tuple(normalized_shape) # type: ignore[arg-type] - self.eps = eps - self.elementwise_affine = elementwise_affine - if self.elementwise_affine: - self.weight = nn.Parameter( - torch.empty(self.normalized_shape, **factory_kwargs) - ) - self.bias = nn.Parameter( - torch.empty(self.normalized_shape, **factory_kwargs) - ) - else: - self.register_parameter("weight", None) - self.register_parameter("bias", None) - - self.reset_parameters() - - def reset_parameters(self) -> None: - if self.elementwise_affine: - nn.init.ones_(self.weight) - nn.init.zeros_(self.bias) - - def forward(self, input: Tensor, embedding: Any = None) -> Tensor: - if isinstance(input, tuple): - input, embedding = input - return ( - F.layer_norm( - input, - self.normalized_shape, - self.weight, - self.bias, - self.eps, - ), - embedding, - ) - - assert embedding is None - return F.layer_norm( - input, self.normalized_shape, self.weight, self.bias, self.eps - ) - - def extra_repr(self) -> str: - return ( - "{normalized_shape}, eps={eps}, " - "elementwise_affine={elementwise_affine}".format(**self.__dict__) - ) - - -class IdentityNorm(nn.Module): - def __init__( - self, - d_model: int, - eps: float = 1e-5, - device=None, - dtype=None, - ) -> None: - super(IdentityNorm, self).__init__() - - def forward(self, input: Tensor, embedding: Any = None) -> Tensor: - if isinstance(input, tuple): - return input - - assert embedding is None - return input - - -class TransformerEncoder(nn.Module): - r"""TransformerEncoder is a stack of N encoder layers. Users can build the - BERT(https://arxiv.org/abs/1810.04805) model with corresponding parameters. - - Args: - encoder_layer: an instance of the TransformerEncoderLayer() class (required). - num_layers: the number of sub-encoder-layers in the encoder (required). - norm: the layer normalization component (optional). - enable_nested_tensor: if True, input will automatically convert to nested tensor - (and convert back on output). This will improve the overall performance of - TransformerEncoder when padding rate is high. Default: ``True`` (enabled). - - Examples:: - >>> encoder_layer = TransformerEncoderLayer(d_model=512, nhead=8) - >>> transformer_encoder = TransformerEncoder(encoder_layer, num_layers=6) - >>> src = torch.rand(10, 32, 512) - >>> out = transformer_encoder(src) - """ - __constants__ = ["norm"] - - def __init__(self, encoder_layer, num_layers, norm=None): - super(TransformerEncoder, self).__init__() - self.layers = _get_clones(encoder_layer, num_layers) - self.num_layers = num_layers - self.norm = norm - - def forward( - self, - src: Tensor, - mask: Optional[Tensor] = None, - src_key_padding_mask: Optional[Tensor] = None, - return_layer_states: bool = False, - cache=None, - ) -> Tensor: - output = src - for mod in self.layers: - output = mod( - output, - src_mask=mask, - src_key_padding_mask=src_key_padding_mask, - cache=cache, - ) - - if self.norm is not None: - output = self.norm(output) - - return output - - -class TransformerEncoderLayer(nn.Module): - __constants__ = ["batch_first", "norm_first"] - def __init__( - self, - d_model: int, - nhead: int, - dim_feedforward: int = 2048, - dropout: float = 0.1, - activation: Union[str, Callable[[Tensor], Tensor]] = F.relu, - batch_first: bool = False, - norm_first: bool = False, - device=None, - dtype=None, - linear1_self_attention_cls: nn.Module = nn.Linear, - linear2_self_attention_cls: nn.Module = nn.Linear, - linear1_feedforward_cls: nn.Module = nn.Linear, - linear2_feedforward_cls: nn.Module = nn.Linear, - layer_norm_cls: nn.Module = LayerNorm, - layer_norm_eps: float = 1e-5, - adaptive_layer_norm=False, - ) -> None: - factory_kwargs = {"device": device, "dtype": dtype} - super(TransformerEncoderLayer, self).__init__() - self.self_attn = MultiheadAttention( - d_model, # 512 16 - nhead, - dropout=dropout, - batch_first=batch_first, - linear1_cls=linear1_self_attention_cls, - linear2_cls=linear2_self_attention_cls, - **factory_kwargs, - ) - self.linear1 = linear1_feedforward_cls( - d_model, dim_feedforward, **factory_kwargs - ) - self.dropout = nn.Dropout(dropout) - self.linear2 = linear2_feedforward_cls( - dim_feedforward, d_model, **factory_kwargs - ) - self.norm_first = norm_first - self.dropout1 = nn.Dropout(dropout) - self.dropout2 = nn.Dropout(dropout) - if isinstance(activation, str): - activation = _get_activation_fn(activation) - elif isinstance(activation, partial): - activation = activation(d_model) - elif activation == BalancedDoubleSwish: - activation = BalancedDoubleSwish(d_model) - self.activation = activation - - norm1 = layer_norm_cls(d_model, eps=layer_norm_eps, **factory_kwargs) - if layer_norm_cls == IdentityNorm: - norm2 = BalancedBasicNorm(d_model, eps=layer_norm_eps, **factory_kwargs) - else: - norm2 = layer_norm_cls(d_model, eps=layer_norm_eps, **factory_kwargs) - - if adaptive_layer_norm: - self.norm1 = AdaptiveLayerNorm(d_model, norm1) - self.norm2 = AdaptiveLayerNorm(d_model, norm2) - else: - self.norm1 = norm1 - self.norm2 = norm2 - - def __setstate__(self, state): - super(TransformerEncoderLayer, self).__setstate__(state) - if not hasattr(self, "activation"): - self.activation = F.relu - - def forward( - self, - src: Tensor, - src_mask: Optional[Tensor] = None, - src_key_padding_mask: Optional[Tensor] = None, - cache=None, - ) -> Tensor: - x = src - stage_embedding = None - x = self.norm1( - x + self._sa_block(x, src_mask, src_key_padding_mask, cache=cache), - stage_embedding, - ) - x = self.norm2(x + self._ff_block(x), stage_embedding) - - return x - - def _sa_block( - self, - x: Tensor, - attn_mask: Optional[Tensor], - key_padding_mask: Optional[Tensor], - cache=None, - ) -> Tensor: - x = self.self_attn( - x, - x, - x, - attn_mask=attn_mask, - key_padding_mask=key_padding_mask, - need_weights=False, - cache=cache, - ) - return self.dropout1(x) - - def _ff_block(self, x: Tensor) -> Tensor: - x = self.linear2(self.dropout(self.activation(self.linear1(x)))) - return self.dropout2(x) - - -class AdaptiveLayerNorm(nn.Module): - r"""Adaptive Layer Normalization""" - - def __init__(self, d_model, norm) -> None: - super(AdaptiveLayerNorm, self).__init__() - self.project_layer = nn.Linear(d_model, 2 * d_model) - self.norm = norm - self.d_model = d_model - self.eps = self.norm.eps - - def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor: - if isinstance(input, tuple): - input, embedding = input - weight, bias = torch.split( - self.project_layer(embedding), - split_size_or_sections=self.d_model, - dim=-1, - ) - return (weight * self.norm(input) + bias, embedding) - - weight, bias = torch.split( - self.project_layer(embedding), - split_size_or_sections=self.d_model, - dim=-1, - ) - return weight * self.norm(input) + bias - - -def _get_clones(module, N): - return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) From f6605299ac360985d30fd3b2cb76a6c0199f43b4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CE=9D=CE=B1=CF=81=CE=BF=CF=85=CF=83=CE=AD=C2=B7=CE=BC?= =?UTF-8?q?=C2=B7=CE=B3=CE=B9=CE=BF=CF=85=CE=BC=CE=B5=CE=BC=CE=AF=C2=B7?= =?UTF-8?q?=CE=A7=CE=B9=CE=BD=CE=B1=CE=BA=CE=AC=CE=BD=CE=BD=CE=B1?= <40709280+NaruseMioShirakana@users.noreply.github.com> Date: Sat, 6 Apr 2024 15:50:38 +0800 Subject: [PATCH 09/14] Add files via upload --- GPT_SoVITS/module/attentions_onnx.py | 354 +++++++++++++++++++++++++++ 1 file changed, 354 insertions(+) create mode 100644 GPT_SoVITS/module/attentions_onnx.py diff --git a/GPT_SoVITS/module/attentions_onnx.py b/GPT_SoVITS/module/attentions_onnx.py new file mode 100644 index 0000000..bc63a06 --- /dev/null +++ b/GPT_SoVITS/module/attentions_onnx.py @@ -0,0 +1,354 @@ +import math +import torch +from torch import nn +from torch.nn import functional as F + +from module import commons +from module.modules import LayerNorm + + +class LayerNorm(nn.Module): + def __init__(self, channels, eps=1e-5): + super().__init__() + self.channels = channels + self.eps = eps + + self.gamma = nn.Parameter(torch.ones(channels)) + self.beta = nn.Parameter(torch.zeros(channels)) + + def forward(self, x): + x = x.transpose(1, -1) + x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps) + return x.transpose(1, -1) + + +@torch.jit.script +def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels): + n_channels_int = n_channels[0] + in_act = input_a + input_b + t_act = torch.tanh(in_act[:, :n_channels_int, :]) + s_act = torch.sigmoid(in_act[:, n_channels_int:, :]) + acts = t_act * s_act + return acts + + +class Encoder(nn.Module): + def __init__( + self, + hidden_channels, + filter_channels, + n_heads, + n_layers, + kernel_size=1, + p_dropout=0.0, + window_size=4, + isflow=True, + **kwargs + ): + super().__init__() + self.hidden_channels = hidden_channels + self.filter_channels = filter_channels + self.n_heads = n_heads + self.n_layers = n_layers + self.kernel_size = kernel_size + self.p_dropout = p_dropout + self.window_size = window_size + # if isflow: + # cond_layer = torch.nn.Conv1d(256, 2*hidden_channels*n_layers, 1) + # self.cond_pre = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, 1) + # self.cond_layer = weight_norm(cond_layer, name='weight') + # self.gin_channels = 256 + self.cond_layer_idx = self.n_layers + if "gin_channels" in kwargs: + self.gin_channels = kwargs["gin_channels"] + if self.gin_channels != 0: + self.spk_emb_linear = nn.Linear(self.gin_channels, self.hidden_channels) + # vits2 says 3rd block, so idx is 2 by default + self.cond_layer_idx = ( + kwargs["cond_layer_idx"] if "cond_layer_idx" in kwargs else 2 + ) + logging.debug(self.gin_channels, self.cond_layer_idx) + assert ( + self.cond_layer_idx < self.n_layers + ), "cond_layer_idx should be less than n_layers" + self.drop = nn.Dropout(p_dropout) + self.attn_layers = nn.ModuleList() + self.norm_layers_1 = nn.ModuleList() + self.ffn_layers = nn.ModuleList() + self.norm_layers_2 = nn.ModuleList() + for i in range(self.n_layers): + self.attn_layers.append( + MultiHeadAttention( + hidden_channels, + hidden_channels, + n_heads, + p_dropout=p_dropout, + window_size=window_size, + ) + ) + self.norm_layers_1.append(LayerNorm(hidden_channels)) + self.ffn_layers.append( + FFN( + hidden_channels, + hidden_channels, + filter_channels, + kernel_size, + p_dropout=p_dropout, + ) + ) + self.norm_layers_2.append(LayerNorm(hidden_channels)) + + def forward(self, x, x_mask, g=None): + attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1) + x = x * x_mask + for i in range(self.n_layers): + if i == self.cond_layer_idx and g is not None: + g = self.spk_emb_linear(g.transpose(1, 2)) + g = g.transpose(1, 2) + x = x + g + x = x * x_mask + y = self.attn_layers[i](x, x, attn_mask) + y = self.drop(y) + x = self.norm_layers_1[i](x + y) + + y = self.ffn_layers[i](x, x_mask) + y = self.drop(y) + x = self.norm_layers_2[i](x + y) + x = x * x_mask + return x + + +class MultiHeadAttention(nn.Module): + def __init__( + self, + channels, + out_channels, + n_heads, + p_dropout=0.0, + window_size=None, + heads_share=True, + block_length=None, + proximal_bias=False, + proximal_init=False, + ): + super().__init__() + assert channels % n_heads == 0 + + self.channels = channels + self.out_channels = out_channels + self.n_heads = n_heads + self.p_dropout = p_dropout + self.window_size = window_size + self.heads_share = heads_share + self.block_length = block_length + self.proximal_bias = proximal_bias + self.proximal_init = proximal_init + self.attn = None + + self.k_channels = channels // n_heads + self.conv_q = nn.Conv1d(channels, channels, 1) + self.conv_k = nn.Conv1d(channels, channels, 1) + self.conv_v = nn.Conv1d(channels, channels, 1) + self.conv_o = nn.Conv1d(channels, out_channels, 1) + self.drop = nn.Dropout(p_dropout) + + if window_size is not None: + n_heads_rel = 1 if heads_share else n_heads + rel_stddev = self.k_channels**-0.5 + self.emb_rel_k = nn.Parameter( + torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) + * rel_stddev + ) + self.emb_rel_v = nn.Parameter( + torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) + * rel_stddev + ) + + nn.init.xavier_uniform_(self.conv_q.weight) + nn.init.xavier_uniform_(self.conv_k.weight) + nn.init.xavier_uniform_(self.conv_v.weight) + if proximal_init: + with torch.no_grad(): + self.conv_k.weight.copy_(self.conv_q.weight) + self.conv_k.bias.copy_(self.conv_q.bias) + + def forward(self, x, c, attn_mask=None): + q = self.conv_q(x) + k = self.conv_k(c) + v = self.conv_v(c) + + x, self.attn = self.attention(q, k, v, mask=attn_mask) + + x = self.conv_o(x) + return x + + def attention(self, query, key, value, mask=None): + # reshape [b, d, t] -> [b, n_h, t, d_k] + b, d, t_s, _ = (*key.size(), query.size(2)) + query = query.view(b, self.n_heads, self.k_channels, -1).transpose(2, 3) + key = key.view(b, self.n_heads, self.k_channels, -1).transpose(2, 3) + value = value.view(b, self.n_heads, self.k_channels, -1).transpose(2, 3) + scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1)) + + if self.window_size is not None: + key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s) + rel_logits = self._matmul_with_relative_keys(query / math.sqrt(self.k_channels), key_relative_embeddings) + scores_local = self._relative_position_to_absolute_position(rel_logits) + scores = scores + scores_local + + if mask is not None: + scores = scores.masked_fill(mask == 0, -1e4) + + p_attn = F.softmax(scores, dim=-1) + p_attn = self.drop(p_attn) + output = torch.matmul(p_attn, value) + + if self.window_size is not None: + relative_weights = self._absolute_position_to_relative_position(p_attn) + value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, t_s) + output = output + self._matmul_with_relative_values(relative_weights, value_relative_embeddings) + + output = (output.transpose(2, 3).contiguous().view(b, d, -1)) + return output, p_attn + + def _matmul_with_relative_values(self, x, y): + """ + x: [b, h, l, m] + y: [h or 1, m, d] + ret: [b, h, l, d] + """ + ret = torch.matmul(x, y.unsqueeze(0)) + return ret + + def _matmul_with_relative_keys(self, x, y): + """ + x: [b, h, l, d] + y: [h or 1, m, d] + ret: [b, h, l, m] + """ + ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1)) + return ret + + def _get_relative_embeddings(self, relative_embeddings, length): + max_relative_position = 2 * self.window_size + 1 + # Pad first before slice to avoid using cond ops. + pad_l = torch.zeros((1), dtype = torch.int64) + length - (self.window_size + 1) + pad_s = torch.zeros((1), dtype = torch.int64) + (self.window_size + 1) - length + pad_length = torch.max(pad_l, other=torch.zeros((1), dtype = torch.int64)) + slice_start_position = torch.max(pad_s, other=torch.zeros((1), dtype = torch.int64)) + + slice_end_position = slice_start_position + 2 * length - 1 + padded_relative_embeddings = F.pad( + relative_embeddings, + commons.convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]), + ) + used_relative_embeddings = padded_relative_embeddings[ + :, slice_start_position:slice_end_position + ] + return used_relative_embeddings + + def _relative_position_to_absolute_position(self, x): + """ + x: [b, h, l, 2*l-1] + ret: [b, h, l, l] + """ + batch, heads, length, _ = x.size() + # Concat columns of pad to shift from relative to absolute indexing. + x = F.pad(x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, 1]])) + + # Concat extra elements so to add up to shape (len+1, 2*len-1). + x_flat = x.view([batch, heads, length * 2 * length]) + x_flat = F.pad( + x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [0, length - 1]]) + ) + + # Reshape and slice out the padded elements. + x_final = x_flat.view([batch, heads, length + 1, 2 * length - 1])[ + :, :, :length, length - 1 : + ] + return x_final + + def _absolute_position_to_relative_position(self, x): + """ + x: [b, h, l, l] + ret: [b, h, l, 2*l-1] + """ + batch, heads, length, _ = x.size() + # padd along column + x = F.pad( + x, commons.convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length - 1]]) + ) + x_flat = x.view([batch, heads, length**2 + length * (length - 1)]) + # add 0's in the beginning that will skew the elements after reshape + x_flat = F.pad(x_flat, commons.convert_pad_shape([[0, 0], [0, 0], [length, 0]])) + x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:] + return x_final + + def _attention_bias_proximal(self, length): + """Bias for self-attention to encourage attention to close positions. + Args: + length: an integer scalar. + Returns: + a Tensor with shape [1, 1, length, length] + """ + r = torch.arange(length, dtype=torch.float32) + diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1) + return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0) + + +class FFN(nn.Module): + def __init__( + self, + in_channels, + out_channels, + filter_channels, + kernel_size, + p_dropout=0.0, + activation=None, + causal=False, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.filter_channels = filter_channels + self.kernel_size = kernel_size + self.p_dropout = p_dropout + self.activation = activation + self.causal = causal + + if causal: + self.padding = self._causal_padding + else: + self.padding = self._same_padding + + self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size) + self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size) + self.drop = nn.Dropout(p_dropout) + + def forward(self, x, x_mask): + x = self.conv_1(self.padding(x * x_mask)) + if self.activation == "gelu": + x = x * torch.sigmoid(1.702 * x) + else: + x = torch.relu(x) + x = self.drop(x) + x = self.conv_2(self.padding(x * x_mask)) + return x * x_mask + + def _causal_padding(self, x): + if self.kernel_size == 1: + return x + pad_l = self.kernel_size - 1 + pad_r = 0 + padding = [[0, 0], [0, 0], [pad_l, pad_r]] + x = F.pad(x, commons.convert_pad_shape(padding)) + return x + + def _same_padding(self, x): + if self.kernel_size == 1: + return x + pad_l = (self.kernel_size - 1) // 2 + pad_r = self.kernel_size // 2 + padding = [[0, 0], [0, 0], [pad_l, pad_r]] + x = F.pad(x, commons.convert_pad_shape(padding)) + return x From 8324bdca23b62e006b79a2440c41bf8b4c024cfb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CE=9D=CE=B1=CF=81=CE=BF=CF=85=CF=83=CE=AD=C2=B7=CE=BC?= =?UTF-8?q?=C2=B7=CE=B3=CE=B9=CE=BF=CF=85=CE=BC=CE=B5=CE=BC=CE=AF=C2=B7?= =?UTF-8?q?=CE=A7=CE=B9=CE=BD=CE=B1=CE=BA=CE=AC=CE=BD=CE=BD=CE=B1?= <40709280+NaruseMioShirakana@users.noreply.github.com> Date: Sat, 6 Apr 2024 15:51:17 +0800 Subject: [PATCH 10/14] Add files via upload From e269b6dc341a195cf4cd43915ae809b7b5ef48c1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CE=9D=CE=B1=CF=81=CE=BF=CF=85=CF=83=CE=AD=C2=B7=CE=BC?= =?UTF-8?q?=C2=B7=CE=B3=CE=B9=CE=BF=CF=85=CE=BC=CE=B5=CE=BC=CE=AF=C2=B7?= =?UTF-8?q?=CE=A7=CE=B9=CE=BD=CE=B1=CE=BA=CE=AC=CE=BD=CE=BD=CE=B1?= <40709280+NaruseMioShirakana@users.noreply.github.com> Date: Thu, 3 Apr 2025 16:40:00 +0800 Subject: [PATCH 11/14] Delete GPT_SoVITS/onnx_export.py --- GPT_SoVITS/onnx_export.py | 344 -------------------------------------- 1 file changed, 344 deletions(-) delete mode 100644 GPT_SoVITS/onnx_export.py diff --git a/GPT_SoVITS/onnx_export.py b/GPT_SoVITS/onnx_export.py deleted file mode 100644 index 43aac19..0000000 --- a/GPT_SoVITS/onnx_export.py +++ /dev/null @@ -1,344 +0,0 @@ -from module.models_onnx import SynthesizerTrn, symbols_v1, symbols_v2 -from AR.models.t2s_lightning_module_onnx import Text2SemanticLightningModule -import torch -import torchaudio -from torch import nn -from feature_extractor import cnhubert - -cnhubert_base_path = "GPT_SoVITS/pretrained_models/chinese-hubert-base" -cnhubert.cnhubert_base_path = cnhubert_base_path -ssl_model = cnhubert.get_model() -from text import cleaned_text_to_sequence -import soundfile -from tools.my_utils import load_audio -import os -import json - -def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False): - hann_window = torch.hann_window(win_size).to( - dtype=y.dtype, device=y.device - ) - y = torch.nn.functional.pad( - y.unsqueeze(1), - (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), - mode="reflect", - ) - y = y.squeeze(1) - spec = torch.stft( - y, - n_fft, - hop_length=hop_size, - win_length=win_size, - window=hann_window, - center=center, - pad_mode="reflect", - normalized=False, - onesided=True, - return_complex=False, - ) - spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6) - return spec - - -class DictToAttrRecursive(dict): - def __init__(self, input_dict): - super().__init__(input_dict) - for key, value in input_dict.items(): - if isinstance(value, dict): - value = DictToAttrRecursive(value) - self[key] = value - setattr(self, key, value) - - def __getattr__(self, item): - try: - return self[item] - except KeyError: - raise AttributeError(f"Attribute {item} not found") - - def __setattr__(self, key, value): - if isinstance(value, dict): - value = DictToAttrRecursive(value) - super(DictToAttrRecursive, self).__setitem__(key, value) - super().__setattr__(key, value) - - def __delattr__(self, item): - try: - del self[item] - except KeyError: - raise AttributeError(f"Attribute {item} not found") - - -class T2SEncoder(nn.Module): - def __init__(self, t2s, vits): - super().__init__() - self.encoder = t2s.onnx_encoder - self.vits = vits - - def forward(self, ref_seq, text_seq, ref_bert, text_bert, ssl_content): - codes = self.vits.extract_latent(ssl_content) - prompt_semantic = codes[0, 0] - bert = torch.cat([ref_bert.transpose(0, 1), text_bert.transpose(0, 1)], 1) - all_phoneme_ids = torch.cat([ref_seq, text_seq], 1) - bert = bert.unsqueeze(0) - prompt = prompt_semantic.unsqueeze(0) - return self.encoder(all_phoneme_ids, bert), prompt - - -class T2SModel(nn.Module): - def __init__(self, t2s_path, vits_model): - super().__init__() - dict_s1 = torch.load(t2s_path, map_location="cpu") - self.config = dict_s1["config"] - self.t2s_model = Text2SemanticLightningModule(self.config, "ojbk", is_train=False) - self.t2s_model.load_state_dict(dict_s1["weight"]) - self.t2s_model.eval() - self.vits_model = vits_model.vq_model - self.hz = 50 - self.max_sec = self.config["data"]["max_sec"] - self.t2s_model.model.top_k = torch.LongTensor([self.config["inference"]["top_k"]]) - self.t2s_model.model.early_stop_num = torch.LongTensor([self.hz * self.max_sec]) - self.t2s_model = self.t2s_model.model - self.t2s_model.init_onnx() - self.onnx_encoder = T2SEncoder(self.t2s_model, self.vits_model) - self.first_stage_decoder = self.t2s_model.first_stage_decoder - self.stage_decoder = self.t2s_model.stage_decoder - #self.t2s_model = torch.jit.script(self.t2s_model) - - def forward(self, ref_seq, text_seq, ref_bert, text_bert, ssl_content): - early_stop_num = self.t2s_model.early_stop_num - - #[1,N] [1,N] [N, 1024] [N, 1024] [1, 768, N] - x, prompts = self.onnx_encoder(ref_seq, text_seq, ref_bert, text_bert, ssl_content) - - prefix_len = prompts.shape[1] - - #[1,N,512] [1,N] - y, k, v, y_emb, x_example = self.first_stage_decoder(x, prompts) - - stop = False - for idx in range(1, 1500): - #[1, N] [N_layer, N, 1, 512] [N_layer, N, 1, 512] [1, N, 512] [1] [1, N, 512] [1, N] - enco = self.stage_decoder(y, k, v, y_emb, x_example) - y, k, v, y_emb, logits, samples = enco - if early_stop_num != -1 and (y.shape[1] - prefix_len) > early_stop_num: - stop = True - if torch.argmax(logits, dim=-1)[0] == self.t2s_model.EOS or samples[0, 0] == self.t2s_model.EOS: - stop = True - if stop: - break - y[0, -1] = 0 - - return y[:, -idx:].unsqueeze(0) - - def export(self, ref_seq, text_seq, ref_bert, text_bert, ssl_content, project_name, dynamo=False): - #self.onnx_encoder = torch.jit.script(self.onnx_encoder) - if dynamo: - export_options = torch.onnx.ExportOptions(dynamic_shapes=True) - onnx_encoder_export_output = torch.onnx.dynamo_export( - self.onnx_encoder, - (ref_seq, text_seq, ref_bert, text_bert, ssl_content), - export_options=export_options - ) - onnx_encoder_export_output.save(f"onnx/{project_name}/{project_name}_t2s_encoder.onnx") - return - - torch.onnx.export( - self.onnx_encoder, - (ref_seq, text_seq, ref_bert, text_bert, ssl_content), - f"onnx/{project_name}/{project_name}_t2s_encoder.onnx", - input_names=["ref_seq", "text_seq", "ref_bert", "text_bert", "ssl_content"], - output_names=["x", "prompts"], - dynamic_axes={ - "ref_seq": {1 : "ref_length"}, - "text_seq": {1 : "text_length"}, - "ref_bert": {0 : "ref_length"}, - "text_bert": {0 : "text_length"}, - "ssl_content": {2 : "ssl_length"}, - }, - opset_version=16 - ) - x, prompts = self.onnx_encoder(ref_seq, text_seq, ref_bert, text_bert, ssl_content) - - torch.onnx.export( - self.first_stage_decoder, - (x, prompts), - f"onnx/{project_name}/{project_name}_t2s_fsdec.onnx", - input_names=["x", "prompts"], - output_names=["y", "k", "v", "y_emb", "x_example"], - dynamic_axes={ - "x": {1 : "x_length"}, - "prompts": {1 : "prompts_length"}, - }, - verbose=False, - opset_version=16 - ) - y, k, v, y_emb, x_example = self.first_stage_decoder(x, prompts) - - torch.onnx.export( - self.stage_decoder, - (y, k, v, y_emb, x_example), - f"onnx/{project_name}/{project_name}_t2s_sdec.onnx", - input_names=["iy", "ik", "iv", "iy_emb", "ix_example"], - output_names=["y", "k", "v", "y_emb", "logits", "samples"], - dynamic_axes={ - "iy": {1 : "iy_length"}, - "ik": {1 : "ik_length"}, - "iv": {1 : "iv_length"}, - "iy_emb": {1 : "iy_emb_length"}, - "ix_example": {1 : "ix_example_length"}, - }, - verbose=False, - opset_version=16 - ) - - -class VitsModel(nn.Module): - def __init__(self, vits_path): - super().__init__() - dict_s2 = torch.load(vits_path,map_location="cpu") - self.hps = dict_s2["config"] - if dict_s2['weight']['enc_p.text_embedding.weight'].shape[0] == 322: - self.hps["model"]["version"] = "v1" - else: - self.hps["model"]["version"] = "v2" - - self.hps = DictToAttrRecursive(self.hps) - self.hps.model.semantic_frame_rate = "25hz" - self.vq_model = SynthesizerTrn( - self.hps.data.filter_length // 2 + 1, - self.hps.train.segment_size // self.hps.data.hop_length, - n_speakers=self.hps.data.n_speakers, - **self.hps.model - ) - self.vq_model.eval() - self.vq_model.load_state_dict(dict_s2["weight"], strict=False) - - def forward(self, text_seq, pred_semantic, ref_audio): - refer = spectrogram_torch( - ref_audio, - self.hps.data.filter_length, - self.hps.data.sampling_rate, - self.hps.data.hop_length, - self.hps.data.win_length, - center=False - ) - return self.vq_model(pred_semantic, text_seq, refer)[0, 0] - - -class GptSoVits(nn.Module): - def __init__(self, vits, t2s): - super().__init__() - self.vits = vits - self.t2s = t2s - - def forward(self, ref_seq, text_seq, ref_bert, text_bert, ref_audio, ssl_content, debug=False): - pred_semantic = self.t2s(ref_seq, text_seq, ref_bert, text_bert, ssl_content) - audio = self.vits(text_seq, pred_semantic, ref_audio) - if debug: - import onnxruntime - sess = onnxruntime.InferenceSession("onnx/koharu/koharu_vits.onnx", providers=["CPU"]) - audio1 = sess.run(None, { - "text_seq" : text_seq.detach().cpu().numpy(), - "pred_semantic" : pred_semantic.detach().cpu().numpy(), - "ref_audio" : ref_audio.detach().cpu().numpy() - }) - return audio, audio1 - return audio - - def export(self, ref_seq, text_seq, ref_bert, text_bert, ref_audio, ssl_content, project_name): - self.t2s.export(ref_seq, text_seq, ref_bert, text_bert, ssl_content, project_name) - pred_semantic = self.t2s(ref_seq, text_seq, ref_bert, text_bert, ssl_content) - torch.onnx.export( - self.vits, - (text_seq, pred_semantic, ref_audio), - f"onnx/{project_name}/{project_name}_vits.onnx", - input_names=["text_seq", "pred_semantic", "ref_audio"], - output_names=["audio"], - dynamic_axes={ - "text_seq": {1 : "text_length"}, - "pred_semantic": {2 : "pred_length"}, - "ref_audio": {1 : "audio_length"}, - }, - opset_version=17, - verbose=False - ) - - -class SSLModel(nn.Module): - def __init__(self): - super().__init__() - self.ssl = ssl_model - - def forward(self, ref_audio_16k): - return self.ssl.model(ref_audio_16k)["last_hidden_state"].transpose(1, 2) - - -def export(vits_path, gpt_path, project_name, vits_model="v2"): - vits = VitsModel(vits_path) - gpt = T2SModel(gpt_path, vits) - gpt_sovits = GptSoVits(vits, gpt) - ssl = SSLModel() - ref_seq = torch.LongTensor([cleaned_text_to_sequence(["n", "i2", "h", "ao3", ",", "w", "o3", "sh", "i4", "b", "ai2", "y", "e4"],version=vits_model)]) - text_seq = torch.LongTensor([cleaned_text_to_sequence(["w", "o3", "sh", "i4", "b", "ai2", "y", "e4", "w", "o3", "sh", "i4", "b", "ai2", "y", "e4", "w", "o3", "sh", "i4", "b", "ai2", "y", "e4"],version=vits_model)]) - ref_bert = torch.randn((ref_seq.shape[1], 1024)).float() - text_bert = torch.randn((text_seq.shape[1], 1024)).float() - ref_audio = torch.randn((1, 48000 * 5)).float() - # ref_audio = torch.tensor([load_audio("rec.wav", 48000)]).float() - ref_audio_16k = torchaudio.functional.resample(ref_audio,48000,16000).float() - ref_audio_sr = torchaudio.functional.resample(ref_audio,48000,vits.hps.data.sampling_rate).float() - - try: - os.mkdir(f"onnx/{project_name}") - except: - pass - - ssl_content = ssl(ref_audio_16k).float() - - # debug = False - debug = True - - # gpt_sovits.export(ref_seq, text_seq, ref_bert, text_bert, ref_audio_sr, ssl_content, project_name) - - if debug: - a, b = gpt_sovits(ref_seq, text_seq, ref_bert, text_bert, ref_audio_sr, ssl_content, debug=debug) - soundfile.write("out1.wav", a.cpu().detach().numpy(), vits.hps.data.sampling_rate) - soundfile.write("out2.wav", b[0], vits.hps.data.sampling_rate) - else: - a = gpt_sovits(ref_seq, text_seq, ref_bert, text_bert, ref_audio_sr, ssl_content).detach().cpu().numpy() - soundfile.write("out.wav", a, vits.hps.data.sampling_rate) - - if vits_model == "v1": - symbols = symbols_v1 - else: - symbols = symbols_v2 - - MoeVSConf = { - "Folder": f"{project_name}", - "Name": f"{project_name}", - "Type": "GPT-SoVits", - "Rate": vits.hps.data.sampling_rate, - "NumLayers": gpt.t2s_model.num_layers, - "EmbeddingDim": gpt.t2s_model.embedding_dim, - "Dict": "BasicDict", - "BertPath": "chinese-roberta-wwm-ext-large", - # "Symbol": symbols, - "AddBlank": False, - } - - MoeVSConfJson = json.dumps(MoeVSConf) - with open(f"onnx/{project_name}.json", 'w') as MoeVsConfFile: - json.dump(MoeVSConf, MoeVsConfFile, indent = 4) - - -if __name__ == "__main__": - try: - os.mkdir("onnx") - except: - pass - - gpt_path = "GPT_weights/nahida-e25.ckpt" - vits_path = "SoVITS_weights/nahida_e30_s3930.pth" - exp_path = "nahida" - export(vits_path, gpt_path, exp_path) - - # soundfile.write("out.wav", a, vits.hps.data.sampling_rate) \ No newline at end of file From e537e129a58eb74da4374083edad2c6de41638b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CE=9D=CE=B1=CF=81=CE=BF=CF=85=CF=83=CE=AD=C2=B7=CE=BC?= =?UTF-8?q?=C2=B7=CE=B3=CE=B9=CE=BF=CF=85=CE=BC=CE=B5=CE=BC=CE=AF=C2=B7?= =?UTF-8?q?=CE=A7=CE=B9=CE=BD=CE=B1=CE=BA=CE=AC=CE=BD=CE=BD=CE=B1?= <40709280+NaruseMioShirakana@users.noreply.github.com> Date: Thu, 3 Apr 2025 16:40:44 +0800 Subject: [PATCH 12/14] onnx export onnx export --- GPT_SoVITS/export_onnx.py | 306 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 306 insertions(+) create mode 100644 GPT_SoVITS/export_onnx.py diff --git a/GPT_SoVITS/export_onnx.py b/GPT_SoVITS/export_onnx.py new file mode 100644 index 0000000..677936c --- /dev/null +++ b/GPT_SoVITS/export_onnx.py @@ -0,0 +1,306 @@ +import os +import json +import onnx +import torch +import onnxsim + +from torch.nn import Module +from feature_extractor import cnhubert +from onnxruntime import InferenceSession +from pytorch_lightning import LightningModule +from transformers import AutoTokenizer, AutoModelForMaskedLM +import AR.models.t2s_model_onnx as t2s + +from module.models_onnx import SynthesizerTrn + +root_path = os.path.dirname(os.path.abspath(__file__)) +onnx_path = os.path.join(root_path, "onnx") +if not os.path.exists(onnx_path): + os.makedirs(onnx_path) + +class BertWrapper(Module): + def __init__(self): + bert_path = os.environ.get( + "bert_path", "GPT_SoVITS/pretrained_models/chinese-roberta-wwm-ext-large" + ) + super(BertWrapper, self).__init__() + self.model = AutoModelForMaskedLM.from_pretrained(bert_path) + self.tokenizer = AutoTokenizer.from_pretrained(bert_path) + + def forward(self, input_ids): + attention_mask = torch.ones_like(input_ids) + token_type_ids = torch.zeros_like(input_ids) + res = self.model(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, output_hidden_states=True) + return torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu()[1:-1] + + def export_onnx(self): + vocab_dict = { k: v for k, v in self.tokenizer.get_vocab().items() } + vocab_path = os.path.join(onnx_path, "Vocab.json") + with open(vocab_path, "w") as f: + json.dump(vocab_dict, f, indent=4) + dummy_input = torch.randint(0, 100, (1, 20)).long() + torch.onnx.export( + self, + dummy_input, + os.path.join(onnx_path, "Bert.onnx"), + input_names=["input_ids"], + output_names=["output"], + dynamic_axes={"input_ids": {0: "batch_size", 1: "sequence_length"}}, + opset_version=18, + ) + sim, _ = onnxsim.simplify(os.path.join(onnx_path, "Bert.onnx")) + onnx.save_model(sim, os.path.join(onnx_path, "Bert.onnx")) + print("Exported BERT to ONNX format.") + + +class CnHubertWrapper(Module): + def __init__(self): + super(CnHubertWrapper, self).__init__() + cnhubert_base_path = os.environ.get( + "cnhubert_base_path", "GPT_SoVITS/pretrained_models/chinese-hubert-base" + ) + cnhubert.cnhubert_base_path = cnhubert_base_path + self.model = cnhubert.get_model().model + + def forward(self, signal): + return self.model(signal)["last_hidden_state"] + + def export_onnx(self): + dummy_input = torch.randn(1, 16000 * 10) + torch.onnx.export( + self, + dummy_input, + os.path.join(onnx_path, "CnHubert.onnx"), + input_names=["signal"], + output_names=["output"], + dynamic_axes={"signal": {0: "batch_size", 1: "sequence_length"}}, + opset_version=18, + ) + sim, _ = onnxsim.simplify(os.path.join(onnx_path, "CnHubert.onnx")) + onnx.save_model(sim, os.path.join(onnx_path, "CnHubert.onnx")) + print("Exported CN-Hubert to ONNX format.") + + +class Text2SemanticLightningModule(LightningModule): + def __init__(self, path, top_k=20, cache_size=2000): + super().__init__() + dict_s1 = torch.load(path, map_location="cpu") + config = dict_s1["config"] + self.model = t2s.Text2SemanticDecoder(config=config) + self.load_state_dict(dict_s1["weight"]) + self.cache_size = cache_size + self.top_k = top_k + +def export_ar(path, top_k=20, cache_size=2000): + model_l = Text2SemanticLightningModule(path, top_k=top_k, cache_size=cache_size) + model = model_l.model + + x = torch.randint(0, 100, (1, 20)).long() + x_len = torch.tensor([20]).long() + y = torch.randint(0, 100, (1, 20)).long() + y_len = torch.tensor([20]).long() + bert_feature = torch.randn(1, 20, 1024) + top_p = torch.tensor([0.8]) + repetition_penalty = torch.tensor([1.35]) + temperature = torch.tensor([0.6]) + + prompt_processor = t2s.PromptProcessor(cache_len=cache_size, model=model, top_k=top_k) + decode_next_token = t2s.DecodeNextToken(cache_len=cache_size, model=model, top_k=top_k) + + torch.onnx.export( + prompt_processor, + (x, x_len, y, y_len, bert_feature, top_p, repetition_penalty, temperature), + os.path.join(onnx_path, "PromptProcessor.onnx"), + input_names=["x", "x_len", "y", "y_len", "bert_feature", "top_p", "repetition_penalty", "temperature"], + output_names=["y", "k_cache", "v_cache", "xy_pos", "y_idx", "samples"], + dynamic_axes={ + "x": {0: "batch_size", 1: "sequence_length"}, + "y": {0: "batch_size", 1: "sequence_length"}, + "bert_feature": {0: "batch_size", 1: "sequence_length"}, + }, + opset_version=18, + ) + + sim, _ = onnxsim.simplify(os.path.join(onnx_path, "PromptProcessor.onnx")) + onnx.save_model(sim, os.path.join(onnx_path, "PromptProcessor.onnx")) + + y, k_cache, v_cache, xy_pos, y_idx, samples = prompt_processor( + x, x_len, y, y_len, bert_feature, top_p, repetition_penalty, temperature + ) + + torch.onnx.export( + decode_next_token, + (y, k_cache, v_cache, xy_pos, y_idx, top_p, repetition_penalty, temperature), + os.path.join(onnx_path, "DecodeNextToken.onnx"), + input_names=["y", "k_cache", "v_cache", "xy_pos", "y_idx", "top_p", "repetition_penalty", "temperature"], + output_names=["y", "k_cache", "v_cache", "xy_pos", "y_idx", "samples"], + dynamic_axes={ + "y": {0: "batch_size", 1: "sequence_length"}, + "k_cache": {1: "batch_size", 2: "sequence_length"}, + "v_cache": {1: "batch_size", 2: "sequence_length"}, + }, + opset_version=18 + ) + + sim, _ = onnxsim.simplify(os.path.join(onnx_path, "DecodeNextToken.onnx")) + onnx.save_model(sim, os.path.join(onnx_path, "DecodeNextToken.onnx")) + + +from io import BytesIO +def load_sovits_new(sovits_path): + f=open(sovits_path,"rb") + meta=f.read(2) + if meta!="PK": + data = b'PK' + f.read() + bio = BytesIO() + bio.write(data) + bio.seek(0) + return torch.load(bio, map_location="cpu", weights_only=False) + return torch.load(sovits_path,map_location="cpu", weights_only=False) + + +class DictToAttrRecursive(dict): + def __init__(self, input_dict): + super().__init__(input_dict) + for key, value in input_dict.items(): + if isinstance(value, dict): + value = DictToAttrRecursive(value) + self[key] = value + setattr(self, key, value) + + def __getattr__(self, item): + try: + return self[item] + except KeyError: + raise AttributeError(f"Attribute {item} not found") + + def __setattr__(self, key, value): + if isinstance(value, dict): + value = DictToAttrRecursive(value) + super(DictToAttrRecursive, self).__setitem__(key, value) + super().__setattr__(key, value) + + def __delattr__(self, item): + try: + del self[item] + except KeyError: + raise AttributeError(f"Attribute {item} not found") + + +def spectrogram_torch(y, n_fft, sampling_rate, hop_size, win_size, center=False): + hann_window = torch.hann_window(win_size).to( + dtype=y.dtype, device=y.device + ) + y = torch.nn.functional.pad( + y.unsqueeze(1), + (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), + mode="reflect", + ) + y = y.squeeze(1) + spec = torch.stft( + y, + n_fft, + hop_length=hop_size, + win_length=win_size, + window=hann_window, + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=False, + ) + spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6) + return spec + + +class Extractor(Module): + def __init__(self, model): + super(Extractor, self).__init__() + self.model = model + + def forward(self, x): + return self.model.extract_latent(x.transpose(1, 2)) + + +class V1V2(Module): + def __init__(self, path): + super(V1V2, self).__init__() + dict_s2 = load_sovits_new(path) + hps = dict_s2["config"] + hps = DictToAttrRecursive(hps) + hps.model.semantic_frame_rate = "25hz" + if 'enc_p.text_embedding.weight'not in dict_s2['weight']: + hps.model.version = "v2"#v3model,v2sybomls + elif dict_s2['weight']['enc_p.text_embedding.weight'].shape[0] == 322: + hps.model.version = "v1" + else: + hps.model.version = "v2" + version=hps.model.version + # print("sovits版本:",hps.model.version) + vq_model = SynthesizerTrn( + hps.data.filter_length // 2 + 1, + hps.train.segment_size // hps.data.hop_length, + n_speakers=hps.data.n_speakers, + **hps.model + ) + vq_model.load_state_dict(dict_s2["weight"], strict=False) + vq_model.eval() + self.vq_model = vq_model + self.hps = hps + self.ext = Extractor(self.vq_model) + + def forward(self, text_seq, pred_semantic, ref_audio): + refer = spectrogram_torch( + ref_audio, + self.hps.data.filter_length, + self.hps.data.sampling_rate, + self.hps.data.hop_length, + self.hps.data.win_length, + center=False + ) + return self.vq_model(pred_semantic.unsqueeze(0), text_seq, refer)[0, 0] + + def export(self): + test_seq = torch.randint(0, 100, (1, 20)).long() + pred_semantic = torch.randint(0, 100, (1, 20)).long() + ref_audio = torch.randn(1, 16000 * 10) + torch.onnx.export( + self, + (test_seq, pred_semantic, ref_audio), + os.path.join(onnx_path, "GptSoVitsV1V2.onnx"), + input_names=["text_seq", "pred_semantic", "ref_audio"], + output_names=["output"], + dynamic_axes={ + "text_seq": {0: "batch_size", 1: "sequence_length"}, + "pred_semantic": {0: "batch_size", 1: "sequence_length"}, + "ref_audio": {0: "batch_size", 1: "sequence_length"}, + }, + opset_version=18, + ) + + sim, _ = onnxsim.simplify(os.path.join(onnx_path, "GptSoVitsV1V2.onnx")) + onnx.save_model(sim, os.path.join(onnx_path, "GptSoVitsV1V2.onnx")) + ref_units = torch.randn(1, 20, 768) + torch.onnx.export( + self.ext, + ref_units, + os.path.join(onnx_path, "Extractor.onnx"), + input_names=["ref_units"], + output_names=["output"], + dynamic_axes={ + "ref_units": {0: "batch_size", 1: "sequence_length"}, + }, + opset_version=18, + ) + + +if __name__ == "__main__": + #CnHubertWrapper().export_onnx() + #BertWrapper().export_onnx() + V1V2("D:\\VSGIT\GPT-SoVITS-main\\GPT_SoVITS\\GPT-SoVITS-v3lora-20250228\\GPT_SoVITS\\t\\SoVITS_weights\\小特.pth").export() + '''export_ar( + "D:\\VSGIT\GPT-SoVITS-main\\GPT_SoVITS\\GPT-SoVITS-v3lora-20250228\\GPT_SoVITS\\t\\GPT_weights\\小特.ckpt", + top_k=10, + cache_size=1500, + )''' + From 70f1ec719ed6dbc91b106f9124d2f8a30b07b3da Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CE=9D=CE=B1=CF=81=CE=BF=CF=85=CF=83=CE=AD=C2=B7=CE=BC?= =?UTF-8?q?=C2=B7=CE=B3=CE=B9=CE=BF=CF=85=CE=BC=CE=B5=CE=BC=CE=AF=C2=B7?= =?UTF-8?q?=CE=A7=CE=B9=CE=BD=CE=B1=CE=BA=CE=AC=CE=BD=CE=BD=CE=B1?= <40709280+NaruseMioShirakana@users.noreply.github.com> Date: Thu, 3 Apr 2025 16:41:13 +0800 Subject: [PATCH 13/14] onnx export onnx export --- GPT_SoVITS/module/models_onnx.py | 345 +++++++++++++++++++++++++------ 1 file changed, 283 insertions(+), 62 deletions(-) diff --git a/GPT_SoVITS/module/models_onnx.py b/GPT_SoVITS/module/models_onnx.py index 1c24056..6bae60c 100644 --- a/GPT_SoVITS/module/models_onnx.py +++ b/GPT_SoVITS/module/models_onnx.py @@ -1,24 +1,28 @@ +import warnings +warnings.filterwarnings("ignore") import copy import math -from typing import Optional +import os +import pdb + import torch from torch import nn from torch.nn import functional as F from module import commons from module import modules -from module import attentions_onnx as attentions - -from f5_tts.model import DiT - +from module import attentions +#from f5_tts.model.backbones.dit import DiT from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm from module.commons import init_weights, get_padding +from module.mrte_model import MRTE from module.quantize import ResidualVectorQuantizer # from text import symbols from text import symbols as symbols_v1 from text import symbols2 as symbols_v2 from torch.cuda.amp import autocast +import contextlib,random class StochasticDurationPredictor(nn.Module): @@ -186,7 +190,7 @@ class TextEncoder(nn.Module): kernel_size, p_dropout, latent_channels=192, - version="v2", + version = "v2", ): super().__init__() self.out_channels = out_channels @@ -220,7 +224,7 @@ class TextEncoder(nn.Module): symbols = symbols_v2.symbols self.text_embedding = nn.Embedding(len(symbols), hidden_channels) - self.mrte = attentions.MRTE() + self.mrte = MRTE() self.encoder2 = attentions.Encoder( hidden_channels, @@ -233,7 +237,7 @@ class TextEncoder(nn.Module): self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1) - def forward(self, y, text, ge, speed=1): + def forward(self, y, text, ge, speed=1,test=None): y_mask = torch.ones_like(y[:1,:1,:]) y = self.ssl_proj(y * y_mask) * y_mask @@ -244,16 +248,35 @@ class TextEncoder(nn.Module): text = self.text_embedding(text).transpose(1, 2) text = self.encoder_text(text * text_mask, text_mask) y = self.mrte(y, y_mask, text, text_mask, ge) - + y = self.encoder2(y * y_mask, y_mask) if(speed!=1): y = F.interpolate(y, size=int(y.shape[-1] / speed)+1, mode="linear") y_mask = F.interpolate(y_mask, size=y.shape[-1], mode="nearest") - + stats = self.proj(y) * y_mask m, logs = torch.split(stats, self.out_channels, dim=1) return y, m, logs, y_mask + def extract_latent(self, x): + x = self.ssl_proj(x) + quantized, codes, commit_loss, quantized_list = self.quantizer(x) + return codes.transpose(0, 1) + + def decode_latent(self, codes, y_mask, refer, refer_mask, ge): + quantized = self.quantizer.decode(codes) + + y = self.vq_proj(quantized) * y_mask + y = self.encoder_ssl(y * y_mask, y_mask) + + y = self.mrte(y, y_mask, refer, refer_mask, ge) + + y = self.encoder2(y * y_mask, y_mask) + + stats = self.proj(y) * y_mask + m, logs = torch.split(stats, self.out_channels, dim=1) + return y, m, logs, y_mask, quantized + class ResidualCouplingBlock(nn.Module): def __init__( @@ -465,7 +488,7 @@ class Generator(torch.nn.Module): if gin_channels != 0: self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1) - def forward(self, x, g:Optional[torch.Tensor]=None): + def forward(self, x, g=None): x = self.conv_pre(x) if g is not None: x = x + self.cond(g) @@ -923,7 +946,7 @@ class SynthesizerTrn(nn.Module): # self.enc_p.encoder_text.requires_grad_(False) # self.enc_p.mrte.requires_grad_(False) - def forward(self, codes, text, refer,noise_scale=0.5, speed=1): + def forward(self, codes, text, refer, noise_scale=0.5): refer_mask = torch.ones_like(refer[:1,:1,:]) if (self.version == "v1"): ge = self.ref_enc(refer * refer_mask, refer_mask) @@ -935,79 +958,98 @@ class SynthesizerTrn(nn.Module): dquantized = torch.cat([quantized, quantized]).permute(1, 2, 0) quantized = dquantized.contiguous().view(1, self.ssl_dim, -1) - x, m_p, logs_p, y_mask = self.enc_p( - quantized, text, ge, speed + _, m_p, logs_p, y_mask = self.enc_p( + quantized, text, ge ) - z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale - z = self.flow(z_p, y_mask, g=ge, reverse=True) - o = self.dec((z * y_mask)[:, :, :], g=ge) return o def extract_latent(self, x): ssl = self.ssl_proj(x) - quantized, codes, commit_loss, quantized_list = self.quantizer(ssl) + _, codes, _, _ = self.quantizer(ssl) return codes.transpose(0, 1) - + + class CFM(torch.nn.Module): def __init__( self, in_channels,dit ): super().__init__() - # self.sigma_min = 1e-6 + self.sigma_min = 1e-6 self.estimator = dit self.in_channels = in_channels - # self.criterion = torch.nn.MSELoss() + self.criterion = torch.nn.MSELoss() - def forward(self, mu:torch.Tensor, x_lens:torch.LongTensor, prompt:torch.Tensor, n_timesteps:torch.LongTensor, temperature:float=1.0): + @torch.inference_mode() + def inference(self, mu, x_lens, prompt, n_timesteps, temperature=1.0, inference_cfg_rate=0): """Forward diffusion""" B, T = mu.size(0), mu.size(1) - x = torch.randn([B, self.in_channels, T], device=mu.device,dtype=mu.dtype) - - ntimesteps = int(n_timesteps) - + x = torch.randn([B, self.in_channels, T], device=mu.device,dtype=mu.dtype) * temperature prompt_len = prompt.size(-1) prompt_x = torch.zeros_like(x,dtype=mu.dtype) prompt_x[..., :prompt_len] = prompt[..., :prompt_len] - x[..., :prompt_len] = 0.0 + x[..., :prompt_len] = 0 mu=mu.transpose(2,1) - t = torch.tensor(0.0,dtype=x.dtype,device=x.device) - d = torch.tensor(1.0/ntimesteps,dtype=x.dtype,device=x.device) - d_tensor = torch.ones(x.shape[0], device=x.device,dtype=mu.dtype) * d - - for j in range(ntimesteps): + t = 0 + d = 1 / n_timesteps + for j in range(n_timesteps): t_tensor = torch.ones(x.shape[0], device=x.device,dtype=mu.dtype) * t - # d_tensor = torch.ones(x.shape[0], device=x.device,dtype=mu.dtype) * d + d_tensor = torch.ones(x.shape[0], device=x.device,dtype=mu.dtype) * d # v_pred = model(x, t_tensor, d_tensor, **extra_args) - v_pred = self.estimator(x, prompt_x, x_lens, t_tensor,d_tensor, mu).transpose(2, 1) - # if inference_cfg_rate>1e-5: - # neg = self.estimator(x, prompt_x, x_lens, t_tensor, d_tensor, mu, use_grad_ckpt=False, drop_audio_cond=True, drop_text=True).transpose(2, 1) - # v_pred=v_pred+(v_pred-neg)*inference_cfg_rate + v_pred = self.estimator(x, prompt_x, x_lens, t_tensor,d_tensor, mu, use_grad_ckpt=False,drop_audio_cond=False,drop_text=False).transpose(2, 1) + if inference_cfg_rate>1e-5: + neg = self.estimator(x, prompt_x, x_lens, t_tensor, d_tensor, mu, use_grad_ckpt=False, drop_audio_cond=True, drop_text=True).transpose(2, 1) + v_pred=v_pred+(v_pred-neg)*inference_cfg_rate x = x + d * v_pred t = t + d - x[:, :, :prompt_len] = 0.0 + x[:, :, :prompt_len] = 0 return x + def forward(self, x1, x_lens, prompt_lens, mu, use_grad_ckpt): + b, _, t = x1.shape + t = torch.rand([b], device=mu.device, dtype=x1.dtype) + x0 = torch.randn_like(x1,device=mu.device) + vt = x1 - x0 + xt = x0 + t[:, None, None] * vt + dt = torch.zeros_like(t,device=mu.device) + prompt = torch.zeros_like(x1) + for i in range(b): + prompt[i, :, :prompt_lens[i]] = x1[i, :, :prompt_lens[i]] + xt[i, :, :prompt_lens[i]] = 0 + gailv=0.3# if ttime()>1736250488 else 0.1 + if random.random() < gailv: + base = torch.randint(2, 8, (t.shape[0],), device=mu.device) + d = 1/torch.pow(2, base) + d_input = d.clone() + d_input[d_input < 1e-2] = 0 + # with torch.no_grad(): + v_pred_1 = self.estimator(xt, prompt, x_lens, t, d_input, mu, use_grad_ckpt).transpose(2, 1).detach() + # v_pred_1 = self.diffusion(xt, t, d_input, cond=conditioning).detach() + x_mid = xt + d[:, None, None] * v_pred_1 + # v_pred_2 = self.diffusion(x_mid, t+d, d_input, cond=conditioning).detach() + v_pred_2 = self.estimator(x_mid, prompt, x_lens, t+d, d_input, mu, use_grad_ckpt).transpose(2, 1).detach() + vt = (v_pred_1 + v_pred_2) / 2 + vt = vt.detach() + dt = 2*d + + vt_pred = self.estimator(xt, prompt, x_lens, t,dt, mu, use_grad_ckpt).transpose(2,1) + loss = 0 + for i in range(b): + loss += self.criterion(vt_pred[i, :, prompt_lens[i]:x_lens[i]], vt[i, :, prompt_lens[i]:x_lens[i]]) + loss /= b + + return loss def set_no_grad(net_g): for name, param in net_g.named_parameters(): param.requires_grad=False -@torch.jit.script_if_tracing -def compile_codes_length(codes): - y_lengths1 = torch.LongTensor([codes.size(2)]).to(codes.device) - return y_lengths1 * 2.5 * 1.5 - -@torch.jit.script_if_tracing -def compile_ref_length(refer): - refer_lengths = torch.LongTensor([refer.size(2)]).to(refer.device) - return refer_lengths class SynthesizerTrnV3(nn.Module): """ @@ -1035,7 +1077,6 @@ class SynthesizerTrnV3(nn.Module): use_sdp=True, semantic_frame_rate=None, freeze_quantizer=None, - version="v3", **kwargs): super().__init__() @@ -1056,7 +1097,6 @@ class SynthesizerTrnV3(nn.Module): self.segment_size = segment_size self.n_speakers = n_speakers self.gin_channels = gin_channels - self.version = version self.model_dim=512 self.use_sdp = use_sdp @@ -1083,7 +1123,7 @@ class SynthesizerTrnV3(nn.Module): n_q=1, bins=1024 ) - freeze_quantizer + self.freeze_quantizer=freeze_quantizer inter_channels2=512 self.bridge=nn.Sequential( nn.Conv1d(inter_channels, inter_channels2, 1, stride=1), @@ -1092,32 +1132,213 @@ class SynthesizerTrnV3(nn.Module): self.wns1=Encoder(inter_channels2, inter_channels2, inter_channels2, 5, 1, 8,gin_channels=gin_channels) self.linear_mel=nn.Conv1d(inter_channels2,100,1,stride=1) self.cfm = CFM(100,DiT(**dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=inter_channels2, conv_layers=4)),)#text_dim is condition feature dim - if freeze_quantizer==True: + if self.freeze_quantizer==True: set_no_grad(self.ssl_proj) set_no_grad(self.quantizer) set_no_grad(self.enc_p) - def create_ge(self, refer): - refer_lengths = compile_ref_length(refer) - refer_mask = torch.unsqueeze(commons.sequence_mask(refer_lengths, refer.size(2)), 1).to(refer.dtype) - ge = self.ref_enc(refer[:,:704] * refer_mask, refer_mask) - return ge + def forward(self, ssl, y, mel,ssl_lengths,y_lengths, text, text_lengths,mel_lengths, use_grad_ckpt):#ssl_lengths no need now + with autocast(enabled=False): + y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, y.size(2)), 1).to(y.dtype) + ge = self.ref_enc(y[:,:704] * y_mask, y_mask) + maybe_no_grad = torch.no_grad() if self.freeze_quantizer else contextlib.nullcontext() + with maybe_no_grad: + if self.freeze_quantizer: + self.ssl_proj.eval()# + self.quantizer.eval() + self.enc_p.eval() + ssl = self.ssl_proj(ssl) + quantized, codes, commit_loss, quantized_list = self.quantizer( + ssl, layers=[0] + ) + quantized = F.interpolate(quantized, scale_factor=2, mode="nearest")##BCT + x, m_p, logs_p, y_mask = self.enc_p(quantized, y_lengths, text, text_lengths, ge) + fea=self.bridge(x) + fea = F.interpolate(fea, scale_factor=1.875, mode="nearest")##BCT + fea, y_mask_ = self.wns1(fea, mel_lengths, ge)##If the 1-minute fine-tuning works fine, no need to manually adjust the learning rate. + B=ssl.shape[0] + prompt_len_max = mel_lengths*2/3 + prompt_len = (torch.rand([B], device=fea.device) * prompt_len_max).floor().to(dtype=torch.long) + minn=min(mel.shape[-1],fea.shape[-1]) + mel=mel[:,:,:minn] + fea=fea[:,:,:minn] + cfm_loss= self.cfm(mel, mel_lengths, prompt_len, fea, use_grad_ckpt) + return cfm_loss - def forward(self, codes, text,ge,speed=1): + @torch.no_grad() + def decode_encp(self, codes,text, refer,ge=None,speed=1): + # print(2333333,refer.shape) + # ge=None + if(ge==None): + refer_lengths = torch.LongTensor([refer.size(2)]).to(refer.device) + refer_mask = torch.unsqueeze(commons.sequence_mask(refer_lengths, refer.size(2)), 1).to(refer.dtype) + ge = self.ref_enc(refer[:,:704] * refer_mask, refer_mask) + y_lengths = torch.LongTensor([int(codes.size(2)*2)]).to(codes.device) + if speed==1: + sizee=int(codes.size(2)*2.5*1.5) + else: + sizee=int(codes.size(2)*2.5*1.5/speed)+1 + y_lengths1 = torch.LongTensor([sizee]).to(codes.device) + text_lengths = torch.LongTensor([text.size(-1)]).to(text.device) - y_lengths1=compile_codes_length(codes) - quantized = self.quantizer.decode(codes) if self.semantic_frame_rate == '25hz': quantized = F.interpolate(quantized, scale_factor=2, mode="nearest")##BCT - x, m_p, logs_p, y_mask = self.enc_p(quantized, text, ge,speed) + x, m_p, logs_p, y_mask = self.enc_p(quantized, y_lengths, text, text_lengths, ge,speed) fea=self.bridge(x) fea = F.interpolate(fea, scale_factor=1.875, mode="nearest")##BCT ####more wn paramter to learn mel fea, y_mask_ = self.wns1(fea, y_lengths1, ge) - return fea + return fea,ge def extract_latent(self, x): ssl = self.ssl_proj(x) quantized, codes, commit_loss, quantized_list = self.quantizer(ssl) - return codes.transpose(0,1) \ No newline at end of file + return codes.transpose(0,1) + + +class SynthesizerTrnV3b(nn.Module): + """ + Synthesizer for Training + """ + + def __init__(self, + spec_channels, + segment_size, + inter_channels, + hidden_channels, + filter_channels, + n_heads, + n_layers, + kernel_size, + p_dropout, + resblock, + resblock_kernel_sizes, + resblock_dilation_sizes, + upsample_rates, + upsample_initial_channel, + upsample_kernel_sizes, + n_speakers=0, + gin_channels=0, + use_sdp=True, + semantic_frame_rate=None, + freeze_quantizer=None, + **kwargs): + + super().__init__() + self.spec_channels = spec_channels + self.inter_channels = inter_channels + self.hidden_channels = hidden_channels + self.filter_channels = filter_channels + self.n_heads = n_heads + self.n_layers = n_layers + self.kernel_size = kernel_size + self.p_dropout = p_dropout + self.resblock = resblock + self.resblock_kernel_sizes = resblock_kernel_sizes + self.resblock_dilation_sizes = resblock_dilation_sizes + self.upsample_rates = upsample_rates + self.upsample_initial_channel = upsample_initial_channel + self.upsample_kernel_sizes = upsample_kernel_sizes + self.segment_size = segment_size + self.n_speakers = n_speakers + self.gin_channels = gin_channels + + self.model_dim=512 + self.use_sdp = use_sdp + self.enc_p = TextEncoder(inter_channels,hidden_channels,filter_channels,n_heads,n_layers,kernel_size,p_dropout) + # self.ref_enc = modules.MelStyleEncoder(spec_channels, style_vector_dim=gin_channels)###Rollback + self.ref_enc = modules.MelStyleEncoder(704, style_vector_dim=gin_channels)###Rollback + self.dec = Generator(inter_channels, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, + upsample_initial_channel, upsample_kernel_sizes, gin_channels=gin_channels) + self.enc_q = PosteriorEncoder(spec_channels, inter_channels, hidden_channels, 5, 1, 16, + gin_channels=gin_channels) + self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels) + + + ssl_dim = 768 + assert semantic_frame_rate in ['25hz', "50hz"] + self.semantic_frame_rate = semantic_frame_rate + if semantic_frame_rate == '25hz': + self.ssl_proj = nn.Conv1d(ssl_dim, ssl_dim, 2, stride=2) + else: + self.ssl_proj = nn.Conv1d(ssl_dim, ssl_dim, 1, stride=1) + + self.quantizer = ResidualVectorQuantizer( + dimension=ssl_dim, + n_q=1, + bins=1024 + ) + self.freeze_quantizer=freeze_quantizer + + inter_channels2=512 + self.bridge=nn.Sequential( + nn.Conv1d(inter_channels, inter_channels2, 1, stride=1), + nn.LeakyReLU() + ) + self.wns1=Encoder(inter_channels2, inter_channels2, inter_channels2, 5, 1, 8,gin_channels=gin_channels) + self.linear_mel=nn.Conv1d(inter_channels2,100,1,stride=1) + self.cfm = CFM(100,DiT(**dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=inter_channels2, conv_layers=4)),)#text_dim is condition feature dim + + + def forward(self, ssl, y, mel,ssl_lengths,y_lengths, text, text_lengths,mel_lengths):#ssl_lengths no need now + with autocast(enabled=False): + y_mask = torch.unsqueeze(commons.sequence_mask(y_lengths, y.size(2)), 1).to(y.dtype) + ge = self.ref_enc(y[:,:704] * y_mask, y_mask) + # ge = self.ref_enc(y * y_mask, y_mask)#change back, new spec setting is whole 24k + # ge=None + maybe_no_grad = torch.no_grad() if self.freeze_quantizer else contextlib.nullcontext() + with maybe_no_grad: + if self.freeze_quantizer: + self.ssl_proj.eval() + self.quantizer.eval() + ssl = self.ssl_proj(ssl) + quantized, codes, commit_loss, quantized_list = self.quantizer( + ssl, layers=[0] + ) + quantized = F.interpolate(quantized, scale_factor=2, mode="nearest")##BCT + x, m_p, logs_p, y_mask = self.enc_p(quantized, y_lengths, text, text_lengths, ge) + z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=ge) + z_p = self.flow(z, y_mask, g=ge) + z_slice, ids_slice = commons.rand_slice_segments(z, y_lengths, self.segment_size) + o = self.dec(z_slice, g=ge) + fea=self.bridge(x) + fea = F.interpolate(fea, scale_factor=1.875, mode="nearest")##BCT + fea, y_mask_ = self.wns1(fea, mel_lengths, ge) + learned_mel = self.linear_mel(fea) + B=ssl.shape[0] + prompt_len_max = mel_lengths*2/3 + prompt_len = (torch.rand([B], device=fea.device) * prompt_len_max).floor().to(dtype=torch.long)# + minn=min(mel.shape[-1],fea.shape[-1]) + mel=mel[:,:,:minn] + fea=fea[:,:,:minn] + cfm_loss= self.cfm(mel, mel_lengths, prompt_len, fea)#fea==cond,y_lengths==target_mel_lengths#ge not need + return commit_loss,cfm_loss,F.mse_loss(learned_mel, mel),o, ids_slice, y_mask, y_mask, (z, z_p, m_p, logs_p, m_q, logs_q), quantized + + @torch.no_grad() + def decode_encp(self, codes,text, refer,ge=None): + # print(2333333,refer.shape) + # ge=None + if(ge==None): + refer_lengths = torch.LongTensor([refer.size(2)]).to(refer.device) + refer_mask = torch.unsqueeze(commons.sequence_mask(refer_lengths, refer.size(2)), 1).to(refer.dtype) + ge = self.ref_enc(refer[:,:704] * refer_mask, refer_mask) + y_lengths = torch.LongTensor([int(codes.size(2)*2)]).to(codes.device) + y_lengths1 = torch.LongTensor([int(codes.size(2)*2.5*1.5)]).to(codes.device) + text_lengths = torch.LongTensor([text.size(-1)]).to(text.device) + + quantized = self.quantizer.decode(codes) + if self.semantic_frame_rate == '25hz': + quantized = F.interpolate(quantized, scale_factor=2, mode="nearest")##BCT + x, m_p, logs_p, y_mask = self.enc_p(quantized, y_lengths, text, text_lengths, ge) + fea=self.bridge(x) + fea = F.interpolate(fea, scale_factor=1.875, mode="nearest")##BCT + ####more wn paramter to learn mel + fea, y_mask_ = self.wns1(fea, y_lengths1, ge) + return fea,ge + + def extract_latent(self, x): + ssl = self.ssl_proj(x) + quantized, codes, commit_loss, quantized_list = self.quantizer(ssl) + return codes.transpose(0,1) + From a889187b84da1ed91b7d99d2a317f930926374f8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CE=9D=CE=B1=CF=81=CE=BF=CF=85=CF=83=CE=AD=C2=B7=CE=BC?= =?UTF-8?q?=C2=B7=CE=B3=CE=B9=CE=BF=CF=85=CE=BC=CE=B5=CE=BC=CE=AF=C2=B7?= =?UTF-8?q?=CE=A7=CE=B9=CE=BD=CE=B1=CE=BA=CE=AC=CE=BD=CE=BD=CE=B1?= <40709280+NaruseMioShirakana@users.noreply.github.com> Date: Thu, 3 Apr 2025 16:42:15 +0800 Subject: [PATCH 14/14] Add files via upload --- GPT_SoVITS/AR/models/t2s_model_onnx.py | 673 ++++++++++++++++++------- 1 file changed, 491 insertions(+), 182 deletions(-) diff --git a/GPT_SoVITS/AR/models/t2s_model_onnx.py b/GPT_SoVITS/AR/models/t2s_model_onnx.py index 92f2d74..51aa35f 100644 --- a/GPT_SoVITS/AR/models/t2s_model_onnx.py +++ b/GPT_SoVITS/AR/models/t2s_model_onnx.py @@ -1,80 +1,88 @@ -# modified from https://github.com/feng-yufei/shared_debugging_code/blob/main/model/t2s_model.py +# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/models/t2s_model.py +# reference: https://github.com/lifeiteng/vall-e +import math +from typing import List, Optional import torch from tqdm import tqdm -from AR.modules.embedding_onnx import SinePositionalEmbedding -from AR.modules.embedding_onnx import TokenEmbedding -from AR.modules.transformer_onnx import LayerNorm -from AR.modules.transformer_onnx import TransformerEncoder -from AR.modules.transformer_onnx import TransformerEncoderLayer +from AR.models.utils import ( + sample, +) +from AR.modules.embedding import TokenEmbedding +from AR.modules.transformer import LayerNorm +from AR.modules.transformer import TransformerEncoder +from AR.modules.transformer import TransformerEncoderLayer from torch import nn from torch.nn import functional as F from torchmetrics.classification import MulticlassAccuracy +from torch.distributions import Exponential + +ISONNXEXPORT = False + default_config = { - "embedding_dim": 512, - "hidden_dim": 512, - "num_head": 8, - "num_layers": 12, - "num_codebook": 8, - "p_dropout": 0.0, - "vocab_size": 1024 + 1, - "phoneme_vocab_size": 512, - "EOS": 1024, + "model": { + "vocab_size": 1025, + "phoneme_vocab_size": 512, + "embedding_dim": 1024, + "hidden_dim": 1024, + "head": 16, + "linear_units": 2048, + "n_layer": 16, + "dropout": 0, + "EOS": 1024, + } } -inf_tensor_value = torch.FloatTensor([-float("Inf")]).float() +def multinomial_sample_one_no_sync( + probs_sort, +): # Does multinomial sampling without a cuda synchronization + q = -torch.log(torch.rand_like(probs_sort)) #https://github.com/RVC-Boss/GPT-SoVITS/pull/835 + return torch.argmax(probs_sort / q, dim=-1, keepdim=True).to(dtype=torch.long) def logits_to_probs( logits, - previous_tokens = None, - temperature: float = 1.0, - top_k = None, - top_p = None, - repetition_penalty: float = 1.0, + previous_tokens: torch.Tensor, + temperature: torch.Tensor, + top_k: torch.Tensor, + top_p: torch.Tensor, + repetition_penalty: torch.Tensor ): - previous_tokens = previous_tokens.squeeze() - if previous_tokens is not None and repetition_penalty != 1.0: - previous_tokens = previous_tokens.long() - score = torch.gather(logits, dim=0, index=previous_tokens) - score = torch.where( - score < 0, score * repetition_penalty, score / repetition_penalty - ) - logits.scatter_(dim=0, index=previous_tokens, src=score) + # if previous_tokens is not None: + # previous_tokens = previous_tokens.squeeze() + # print(logits.shape,previous_tokens.shape) + # pdb.set_trace() + previous_tokens = previous_tokens.long() + score = torch.gather(logits, dim=1, index=previous_tokens) + score = torch.where( + score < 0, score * repetition_penalty, score / repetition_penalty + ) + logits.scatter_(dim=1, index=previous_tokens, src=score) - if top_p is not None and top_p < 1.0: - sorted_logits, sorted_indices = torch.sort(logits, descending=True) - cum_probs = torch.cumsum( - torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1 - ) - sorted_indices_to_remove = cum_probs > top_p - sorted_indices_to_remove[0] = False # keep at least one option - indices_to_remove = sorted_indices_to_remove.scatter( - dim=0, index=sorted_indices, src=sorted_indices_to_remove - ) - logits = logits.masked_fill(indices_to_remove, -float("Inf")) + sorted_logits, sorted_indices = torch.sort(logits, descending=True) + cum_probs = torch.cumsum( + torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1 + ) + sorted_indices_to_remove = cum_probs > top_p + sorted_indices_to_remove[:, 0] = False # keep at least one option + indices_to_remove = sorted_indices_to_remove.scatter( + dim=1, index=sorted_indices, src=sorted_indices_to_remove + ) + logits = logits.masked_fill(indices_to_remove, -float("Inf")) - logits = logits / max(temperature, 1e-5) + logits = logits / torch.clamp_min(temperature, 1e-5) - if top_k is not None: - v, _ = torch.topk(logits, top_k) - pivot = v.select(-1, -1).unsqueeze(-1) - logits = torch.where(logits < pivot, inf_tensor_value, logits) + v, _ = torch.topk(logits, top_k) + pivot = v[: , -1].unsqueeze(-1) + logits = torch.where(logits < pivot, -float("Inf"), logits) probs = torch.nn.functional.softmax(logits, dim=-1) return probs -def multinomial_sample_one_no_sync( - probs_sort -): # Does multinomial sampling without a cuda synchronization - q = torch.randn_like(probs_sort) - return torch.argmax(probs_sort / q, dim=-1, keepdim=True).to(dtype=torch.int) - - def sample( logits, - previous_tokens, + previous_tokens: Optional[torch.Tensor] = None, **sampling_kwargs, ): probs = logits_to_probs( @@ -83,125 +91,326 @@ def sample( idx_next = multinomial_sample_one_no_sync(probs) return idx_next, probs +# @torch.jit.script ## 使用的话首次推理会非常慢，而且推理速度不稳定 +# Efficient implementation equivalent to the following: +def scaled_dot_product_attention(query:torch.Tensor, key:torch.Tensor, value:torch.Tensor, attn_mask:Optional[torch.Tensor]=None, scale:Optional[torch.Tensor]=None) -> torch.Tensor: + B, H, L, S =query.size(0), query.size(1), query.size(-2), key.size(-2) + if scale is None: + scale_factor = torch.tensor(1 / math.sqrt(query.size(-1))) + else: + scale_factor = scale + attn_bias = torch.zeros(B, H, L, S, dtype=query.dtype, device=query.device) -class OnnxEncoder(nn.Module): - def __init__(self, ar_text_embedding, bert_proj, ar_text_position): - super().__init__() - self.ar_text_embedding = ar_text_embedding - self.bert_proj = bert_proj - self.ar_text_position = ar_text_position + if attn_mask is not None: + if attn_mask.dtype == torch.bool: + attn_bias = attn_bias.masked_fill(attn_mask, float("-inf")) + else: + attn_bias = attn_bias + attn_mask + attn_weight = query @ key.transpose(-2, -1) * scale_factor + attn_weight += attn_bias + attn_weight = torch.softmax(attn_weight, dim=-1) + + if attn_mask is not None: + if attn_mask.dtype == torch.bool: + attn_weight = attn_weight.masked_fill(attn_mask, 0) + else: + attn_mask = attn_mask.clone() + attn_mask[attn_mask!=float("-inf")] =0 + attn_mask[attn_mask==float("-inf")] =1 + attn_weight = attn_weight.masked_fill(attn_mask, 0) + + return attn_weight @ value + +@torch.jit.script +class T2SMLP: + def __init__(self, w1, b1, w2, b2): + self.w1 = w1 + self.b1 = b1 + self.w2 = w2 + self.b2 = b2 + + def forward(self, x): + x = F.relu(F.linear(x, self.w1, self.b1)) + x = F.linear(x, self.w2, self.b2) + return x + + +@torch.jit.script +class T2SBlock: + def __init__( + self, + num_heads, + hidden_dim: int, + mlp: T2SMLP, + qkv_w, + qkv_b, + out_w, + out_b, + norm_w1, + norm_b1, + norm_eps1, + norm_w2, + norm_b2, + norm_eps2, + ): + self.num_heads = num_heads + self.mlp = mlp + self.hidden_dim: int = hidden_dim + self.qkv_w = qkv_w + self.qkv_b = qkv_b + self.out_w = out_w + self.out_b = out_b + self.norm_w1 = norm_w1 + self.norm_b1 = norm_b1 + self.norm_eps1 = norm_eps1 + self.norm_w2 = norm_w2 + self.norm_b2 = norm_b2 + self.norm_eps2 = norm_eps2 + + self.false = torch.tensor(False, dtype=torch.bool) + + @torch.jit.ignore + def to_mask(self, x:torch.Tensor, padding_mask:Optional[torch.Tensor]): + if padding_mask is None: + return x + + if padding_mask.dtype == torch.bool: + return x.masked_fill(padding_mask, 0) + else: + return x * padding_mask + + def process_prompt(self, x:torch.Tensor, attn_mask : torch.Tensor, padding_mask:Optional[torch.Tensor]=None, torch_sdpa:bool=True): + q, k, v = F.linear(self.to_mask(x, padding_mask), self.qkv_w, self.qkv_b).chunk(3, dim=-1) + + batch_size = q.shape[0] + q_len = q.shape[1] + kv_len = k.shape[1] + + q = self.to_mask(q, padding_mask) + k_cache = self.to_mask(k, padding_mask) + v_cache = self.to_mask(v, padding_mask) + + q = q.view(batch_size, q_len, self.num_heads, -1).transpose(1, 2) + k = k_cache.view(batch_size, kv_len, self.num_heads, -1).transpose(1, 2) + v = v_cache.view(batch_size, kv_len, self.num_heads, -1).transpose(1, 2) + + if torch_sdpa: + attn = F.scaled_dot_product_attention(q, k, v, ~attn_mask) + else: + attn = scaled_dot_product_attention(q, k, v, attn_mask) + + attn = attn.transpose(1, 2).reshape(batch_size, q_len, -1) + attn = F.linear(self.to_mask(attn, padding_mask), self.out_w, self.out_b) + + x = x + attn + x = F.layer_norm( + x, [self.hidden_dim], self.norm_w1, self.norm_b1, self.norm_eps1 + ) + x = x + self.mlp.forward(x) + x = F.layer_norm( + x, + [self.hidden_dim], + self.norm_w2, + self.norm_b2, + self.norm_eps2, + ) + return x, k_cache, v_cache - def forward(self, x, bert_feature): - x = self.ar_text_embedding(x) - x = x + self.bert_proj(bert_feature.transpose(1, 2)) - return self.ar_text_position(x) + def decode_next_token(self, x, k_cache, v_cache): + q, k, v = F.linear(x, self.qkv_w, self.qkv_b).chunk(3, dim=-1) + + k_cache = torch.cat([k_cache, k], dim=1) + v_cache = torch.cat([v_cache, v], dim=1) + + batch_size = q.shape[0] + q_len = q.shape[1] + kv_len = k_cache.shape[1] + + q = q.view(batch_size, q_len, self.num_heads, -1).transpose(1, 2) + k = k_cache.view(batch_size, kv_len, self.num_heads, -1).transpose(1, 2) + v = v_cache.view(batch_size, kv_len, self.num_heads, -1).transpose(1, 2) + + attn = F.scaled_dot_product_attention(q, k, v) + + attn = attn.transpose(1, 2).reshape(batch_size, q_len, -1) + attn = F.linear(attn, self.out_w, self.out_b) + + x = x + attn + x = F.layer_norm( + x, [self.hidden_dim], self.norm_w1, self.norm_b1, self.norm_eps1 + ) + x = x + self.mlp.forward(x) + x = F.layer_norm( + x, + [self.hidden_dim], + self.norm_w2, + self.norm_b2, + self.norm_eps2, + ) + return x, k_cache, v_cache -class T2SFirstStageDecoder(nn.Module): - def __init__(self, ar_audio_embedding, ar_audio_position, h, ar_predict_layer, loss_fct, ar_accuracy_metric, - top_k, early_stop_num, num_layers): +@torch.jit.script +class T2STransformer: + def __init__(self, num_blocks : int, blocks: List[T2SBlock]): + self.num_blocks : int = num_blocks + self.blocks = blocks + + def process_prompt( + self, x:torch.Tensor, attn_mask : torch.Tensor, + padding_mask : Optional[torch.Tensor]=None, + torch_sdpa:bool=True + ): + k_cache : List[torch.Tensor] = [] + v_cache : List[torch.Tensor] = [] + for i in range(self.num_blocks): + x, k_cache_, v_cache_ = self.blocks[i].process_prompt(x, attn_mask, padding_mask, torch_sdpa) + k_cache.append(k_cache_) + v_cache.append(v_cache_) + return x, k_cache, v_cache + + def decode_next_token( + self, x:torch.Tensor, + k_cache, + v_cache, + ): + K_Cache = [] + V_Cache = [] + for i in range(self.num_blocks): + x, k, v = self.blocks[i].decode_next_token(x, k_cache[i], v_cache[i]) + K_Cache.append(k) + V_Cache.append(v) + K_Cache = torch.stack(K_Cache, dim=0) + V_Cache = torch.stack(V_Cache, dim=0) + return x, K_Cache, V_Cache + + +class SinePositionalEmbedding(nn.Module): + def __init__( + self, + embedding_dim: int, + dropout: float = 0.0, + scale: bool = False, + alpha: bool = False, + ): super().__init__() - self.ar_audio_embedding = ar_audio_embedding - self.ar_audio_position = ar_audio_position - self.h = h - self.ar_predict_layer = ar_predict_layer - self.loss_fct = loss_fct - self.ar_accuracy_metric = ar_accuracy_metric + self.embedding_dim = embedding_dim + self.x_scale = math.sqrt(embedding_dim) if scale else 1.0 + self.alpha = nn.Parameter(torch.ones(1), requires_grad=alpha) + self.dropout = torch.nn.Dropout(p=dropout) + + self.reverse = False + self.pe = None + self.extend_pe(torch.tensor(0.0).expand(1, 114514)) + + def extend_pe(self, x): + if self.pe is not None: + if self.pe.size(1) >= x.size(1): + if self.pe.dtype != x.dtype or self.pe.device != x.device: + self.pe = self.pe.to(dtype=x.dtype, device=x.device) + return + pe = torch.zeros(x.size(1), self.embedding_dim) + if self.reverse: + position = torch.arange( + x.size(1) - 1, -1, -1.0, dtype=torch.float32 + ).unsqueeze(1) + else: + position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, self.embedding_dim, 2, dtype=torch.float32) + * -(math.log(10000.0) / self.embedding_dim) + ) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + pe = pe.unsqueeze(0) + self.pe = pe.to(device=x.device, dtype=x.dtype).detach() + + def forward(self, x: torch.Tensor, x_size) -> torch.Tensor: + output = x.unsqueeze(-1) if x.ndim == 2 else x + output[:,:x_size,:] = output[:,:x_size,:] * self.x_scale + self.alpha * self.pe[:, : x_size] + return self.dropout(output) + + +class PromptProcessor(nn.Module): + def __init__(self, cache_len, model, top_k): + super(PromptProcessor, self).__init__() self.top_k = top_k - self.early_stop_num = early_stop_num - self.num_layers = num_layers - - def forward(self, x, prompt): - y = prompt - x_example = x[:,:,0] * 0.0 - #N, 1, 512 - cache = { - "all_stage": self.num_layers, - "k": None, - "v": None, - "y_emb": None, - "first_infer": 1, - "stage": 0, - } + self.model = model + self.ar_text_embedding = model.ar_text_embedding + self.ar_text_position = model.ar_text_position + self.ar_audio_embedding = model.ar_audio_embedding + self.ar_audio_position = model.ar_audio_position + self.bert_proj = model.bert_proj + cache_len = torch.tensor(cache_len) + self.register_buffer("cache_len", cache_len, persistent=False) + + def forward(self, x, x_len, y, y_len, bert_feature, top_p, repetition_penalty, temperature): + bsz = x.size(0) + src_len = x_len + y_len + + x_emb = self.ar_text_embedding(x) + x_emb = x_emb + self.bert_proj(bert_feature) + x_pos = self.ar_text_position(x_emb, x_len) + x_attn_mask = torch.zeros((x_len, x_len), dtype=torch.bool) y_emb = self.ar_audio_embedding(y) + y_pos = self.ar_audio_position(y_emb, y_len) + y_attn_mask = F.pad(torch.triu(torch.ones(y_len, y_len, dtype=torch.bool), diagonal=1),(x_len, 0),value=False) - cache["y_emb"] = y_emb - y_pos = self.ar_audio_position(y_emb) + xy_pos = torch.concat([x_pos, y_pos], dim=1) - xy_pos = torch.concat([x, y_pos], dim=1) + x_attn_mask_pad = F.pad(x_attn_mask,(0, y_len),value=True) + xy_attn_mask = torch.concat([x_attn_mask_pad, y_attn_mask], dim=0).unsqueeze(0)\ + .expand(bsz * self.model.num_head, -1, -1)\ + .view(bsz, self.model.num_head, src_len, src_len)\ + .to(device=x.device, dtype=torch.bool) - y_example = y_pos[:,:,0] * 0.0 - x_attn_mask = torch.matmul(x_example.transpose(0, 1) , x_example).bool() - y_attn_mask = torch.ones_like(torch.matmul(y_example.transpose(0, 1), y_example), dtype=torch.int64) - y_attn_mask = torch.cumsum(y_attn_mask, dim=1) - torch.cumsum( - torch.ones_like(y_example.transpose(0, 1), dtype=torch.int64), dim=0 - ) - y_attn_mask = y_attn_mask > 0 - - x_y_pad = torch.matmul(x_example.transpose(0, 1), y_example).bool() - y_x_pad = torch.matmul(y_example.transpose(0, 1), x_example).bool() - x_attn_mask_pad = torch.cat([x_attn_mask, torch.ones_like(x_y_pad)], dim=1) - y_attn_mask = torch.cat([y_x_pad, y_attn_mask], dim=1) - xy_attn_mask = torch.concat([x_attn_mask_pad, y_attn_mask], dim=0) - cache["k"] = torch.matmul(x_attn_mask_pad[0].float().unsqueeze(-1), torch.zeros((1, 512)))\ - .unsqueeze(1).repeat(self.num_layers, 1, 1, 1) - cache["v"] = torch.matmul(x_attn_mask_pad[0].float().unsqueeze(-1), torch.zeros((1, 512)))\ - .unsqueeze(1).repeat(self.num_layers, 1, 1, 1) - - xy_dec = self.h(xy_pos, mask=xy_attn_mask, cache=cache) - logits = self.ar_predict_layer(xy_dec[:, -1]) - samples = sample(logits[0], y, top_k=self.top_k, top_p=1.0, repetition_penalty=1.35)[0].unsqueeze(0) - - y = torch.concat([y, samples], dim=1) - - return y, cache["k"], cache["v"], cache["y_emb"], x_example - - -class T2SStageDecoder(nn.Module): - def __init__(self, ar_audio_embedding, ar_audio_position, h, ar_predict_layer, loss_fct, ar_accuracy_metric, - top_k, early_stop_num, num_layers): - super().__init__() - self.ar_audio_embedding = ar_audio_embedding - self.ar_audio_position = ar_audio_position - self.h = h - self.ar_predict_layer = ar_predict_layer - self.loss_fct = loss_fct - self.ar_accuracy_metric = ar_accuracy_metric - self.top_k = top_k - self.early_stop_num = early_stop_num - self.num_layers = num_layers - - def forward(self, y, k, v, y_emb, x_example): - cache = { - "all_stage": self.num_layers, - "k": torch.nn.functional.pad(k, (0, 0, 0, 0, 0, 1)), - "v": torch.nn.functional.pad(v, (0, 0, 0, 0, 0, 1)), - "y_emb": y_emb, - "first_infer": 0, - "stage": 0, - } - - y_emb = torch.cat( - [cache["y_emb"], self.ar_audio_embedding(y[:, -1:])], 1 - ) - cache["y_emb"] = y_emb - y_pos = self.ar_audio_position(y_emb) - - xy_pos = y_pos[:, -1:] + xy_dec, k_cache, v_cache = self.model.t2s_transformer.process_prompt(xy_pos, xy_attn_mask, None) - y_example = y_pos[:,:,0] * 0.0 - - xy_attn_mask = torch.cat([x_example, y_example], dim=1) - xy_attn_mask = torch.zeros_like(xy_attn_mask, dtype=torch.bool) - - xy_dec = self.h(xy_pos, mask=xy_attn_mask, cache=cache) - logits = self.ar_predict_layer(xy_dec[:, -1]) - samples = sample(logits[0], y, top_k=self.top_k, top_p=1.0, repetition_penalty=1.35)[0].unsqueeze(0) + logits = self.model.ar_predict_layer( + xy_dec[:, -1] + ) + samples = sample( + logits, y, top_k=self.top_k, top_p=top_p, repetition_penalty=repetition_penalty, temperature=temperature + )[0] y = torch.concat([y, samples], dim=1) - return y, cache["k"], cache["v"], cache["y_emb"], logits, samples + y_emb = self.ar_audio_embedding(samples) + xy_pos = y_emb * self.ar_audio_position.x_scale + self.ar_audio_position.alpha * self.ar_audio_position.pe[:, y_len].to(dtype=y_emb.dtype,device=y_emb.device) + k_cache = torch.stack(k_cache, dim=0) + v_cache = torch.stack(v_cache, dim=0) + return y, k_cache, v_cache, xy_pos, y_len + 1, samples + + +class DecodeNextToken(nn.Module): + def __init__(self, cache_len, model, top_k): + super(DecodeNextToken, self).__init__() + self.top_k = top_k + self.model = model + self.ar_text_embedding = model.ar_text_embedding + self.ar_text_position = model.ar_text_position + self.ar_audio_embedding = model.ar_audio_embedding + self.ar_audio_position = model.ar_audio_position + cache_len = torch.tensor(cache_len) + self.register_buffer("cache_len", cache_len, persistent=False) + + def forward(self, y, k_cache, v_cache, xy_pos, y_idx, top_p, repetition_penalty, temperature): + xy_dec, k_cache, v_cache = self.model.t2s_transformer.decode_next_token(xy_pos, k_cache, v_cache) + logits = self.model.ar_predict_layer( + xy_dec[:, -1] + ) + + samples = sample( + logits, y, top_k=self.top_k, top_p=top_p, repetition_penalty=repetition_penalty, temperature=temperature + )[0] + y = torch.concat([y, samples], dim=1) + + y_emb = self.ar_audio_embedding(samples) + xy_pos = y_emb * self.ar_audio_position.x_scale + self.ar_audio_position.alpha * self.ar_audio_position.pe[:, y_idx].to(dtype=y_emb.dtype,device=y_emb.device) + + return y, k_cache, v_cache, xy_pos, y_idx + 1, samples + class Text2SemanticDecoder(nn.Module): @@ -214,15 +423,24 @@ class Text2SemanticDecoder(nn.Module): self.norm_first = norm_first self.vocab_size = config["model"]["vocab_size"] self.phoneme_vocab_size = config["model"]["phoneme_vocab_size"] - self.p_dropout = float(config["model"]["dropout"]) + self.p_dropout = config["model"]["dropout"] self.EOS = config["model"]["EOS"] self.norm_first = norm_first assert self.EOS == self.vocab_size - 1 self.bert_proj = nn.Linear(1024, self.embedding_dim) - self.ar_text_embedding = TokenEmbedding(self.embedding_dim, self.phoneme_vocab_size, self.p_dropout) - self.ar_text_position = SinePositionalEmbedding(self.embedding_dim, dropout=0.1, scale=False, alpha=True) - self.ar_audio_embedding = TokenEmbedding(self.embedding_dim, self.vocab_size, self.p_dropout) - self.ar_audio_position = SinePositionalEmbedding(self.embedding_dim, dropout=0.1, scale=False, alpha=True) + self.ar_text_embedding = TokenEmbedding( + self.embedding_dim, self.phoneme_vocab_size, self.p_dropout + ) + self.ar_text_position = SinePositionalEmbedding( + self.embedding_dim, dropout=0.1, scale=False, alpha=True + ) + self.ar_audio_embedding = TokenEmbedding( + self.embedding_dim, self.vocab_size, self.p_dropout + ) + self.ar_audio_position = SinePositionalEmbedding( + self.embedding_dim, dropout=0.1, scale=False, alpha=True + ) + self.h = TransformerEncoder( TransformerEncoderLayer( d_model=self.model_dim, @@ -235,8 +453,10 @@ class Text2SemanticDecoder(nn.Module): num_layers=self.num_layers, norm=LayerNorm(self.model_dim) if norm_first else None, ) + self.ar_predict_layer = nn.Linear(self.model_dim, self.vocab_size, bias=False) self.loss_fct = nn.CrossEntropyLoss(reduction="sum") + self.ar_accuracy_metric = MulticlassAccuracy( self.vocab_size, top_k=top_k, @@ -244,38 +464,127 @@ class Text2SemanticDecoder(nn.Module): multidim_average="global", ignore_index=self.EOS, ) - self.top_k = torch.LongTensor([1]) - self.early_stop_num = torch.LongTensor([-1]) - def init_onnx(self): - self.onnx_encoder = OnnxEncoder(self.ar_text_embedding, self.bert_proj, self.ar_text_position) - self.first_stage_decoder = T2SFirstStageDecoder(self.ar_audio_embedding, self.ar_audio_position, self.h, - self.ar_predict_layer, self.loss_fct, self.ar_accuracy_metric, self.top_k, self.early_stop_num, - self.num_layers) - self.stage_decoder = T2SStageDecoder(self.ar_audio_embedding, self.ar_audio_position, self.h, - self.ar_predict_layer, self.loss_fct, self.ar_accuracy_metric, self.top_k, self.early_stop_num, - self.num_layers) + blocks = [] - def forward(self, x, prompts, bert_feature): - early_stop_num = self.early_stop_num - prefix_len = prompts.shape[1] + for i in range(self.num_layers): + layer = self.h.layers[i] + t2smlp = T2SMLP( + layer.linear1.weight, + layer.linear1.bias, + layer.linear2.weight, + layer.linear2.bias + ) - x = self.onnx_encoder(x, bert_feature) - y, k, v, y_emb, stage, x_example = self.first_stage_decoder(x, prompts) + block = T2SBlock( + self.num_head, + self.model_dim, + t2smlp, + layer.self_attn.in_proj_weight, + layer.self_attn.in_proj_bias, + layer.self_attn.out_proj.weight, + layer.self_attn.out_proj.bias, + layer.norm1.weight, + layer.norm1.bias, + layer.norm1.eps, + layer.norm2.weight, + layer.norm2.bias, + layer.norm2.eps + ) + blocks.append(block) + + self.t2s_transformer = T2STransformer(self.num_layers, blocks) + + def infer_panel_naive( + self, + x:torch.LongTensor, #####全部文本token + x_lens:torch.LongTensor, + prompts:torch.LongTensor, ####参考音频token + bert_feature:torch.LongTensor, + top_k: int = -100, + top_p: int = 100, + early_stop_num: int = -1, + temperature: float = 1.0, + repetition_penalty: float = 1.35, + **kwargs + ): + x = self.ar_text_embedding(x) + x = x + self.bert_proj(bert_feature.transpose(1, 2)) + x = self.ar_text_position(x) + + y = prompts + + x_len = x.shape[1] + x_attn_mask = torch.zeros((x_len, x_len), dtype=torch.bool) stop = False - for idx in range(1, 1500): - enco = self.stage_decoder(y, k, v, y_emb, stage, x_example) - y, k, v, y_emb, stage, logits, samples = enco + k_cache = None + v_cache = None + + y_emb = self.ar_audio_embedding(y) + y_len = y_emb.shape[1] + prefix_len = y.shape[1] + y_pos = self.ar_audio_position(y_emb) + xy_pos = torch.concat([x, y_pos], dim=1) + + bsz = x.shape[0] + src_len = x_len + y_len + x_attn_mask_pad = F.pad( + x_attn_mask, + (0, y_len), + value=True, + ) + y_attn_mask = F.pad( + torch.triu(torch.ones(y_len, y_len, dtype=torch.bool), diagonal=1), + (x_len, 0), + value=False, + ) + xy_attn_mask = torch.concat([x_attn_mask_pad, y_attn_mask], dim=0)\ + .unsqueeze(0)\ + .expand(bsz*self.num_head, -1, -1)\ + .view(bsz, self.num_head, src_len, src_len)\ + .to(device=x.device, dtype=torch.bool) + + for idx in tqdm(range(1500)): + if xy_attn_mask is not None: + xy_dec, k_cache, v_cache = self.t2s_transformer.process_prompt(xy_pos, xy_attn_mask, None) + else: + xy_dec, k_cache, v_cache = self.t2s_transformer.decode_next_token(xy_pos, k_cache, v_cache) + + logits = self.ar_predict_layer( + xy_dec[:, -1] + ) + + if idx == 0: + xy_attn_mask = None + if(idx<11):###至少预测出10个token不然不给停止（0.4s） + logits = logits[:, :-1] + + samples = sample( + logits, y, top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty, temperature=temperature + )[0] + + y = torch.concat([y, samples], dim=1) + if early_stop_num != -1 and (y.shape[1] - prefix_len) > early_stop_num: + print("use early stop num:", early_stop_num) stop = True + if torch.argmax(logits, dim=-1)[0] == self.EOS or samples[0, 0] == self.EOS: stop = True if stop: + if y.shape[1] == 0: + y = torch.concat([y, torch.zeros_like(samples)], dim=1) + print("bad zero prediction") + print(f"T2S Decoding EOS [{prefix_len} -> {y.shape[1]}]") break - y[0, -1] = 0 - return y, idx + ####################### update next step ################################### + y_emb = self.ar_audio_embedding(y[:, -1:]) + xy_pos = y_emb * self.ar_audio_position.x_scale + self.ar_audio_position.alpha * self.ar_audio_position.pe[:, y_len + idx].to(dtype=y_emb.dtype,device=y_emb.device) + + return y[:, :-1], idx + def infer(self, x, prompts, bert_feature): top_k = self.top_k early_stop_num = self.early_stop_num