xinference 1.4.1__py3-none-any.whl → 1.5.0.post1__py3-none-any.whl

xinference/thirdparty/megatts3/tts/modules/ar_dur/ar_dur_predictor.py

@@ -0,0 +1,362 @@
+# Copyright 2025 ByteDance and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import random
+from copy import deepcopy
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn import Linear
+from tqdm import tqdm
+
+from tts.modules.ar_dur.commons.layers import Embedding, LayerNorm
+from tts.modules.ar_dur.commons.nar_tts_modules import PosEmb
+from tts.modules.ar_dur.commons.rot_transformer import RotTransformerDecoderLayer
+from tts.modules.ar_dur.commons.transformer import SinusoidalPositionalEmbedding
+from tts.modules.ar_dur.commons.rel_transformer import RelTransformerEncoder
+
+FS_ENCODERS = {
+    'rel_fft': lambda hp, dict_size: RelTransformerEncoder(
+        dict_size, hp['hidden_size'], hp['hidden_size'],
+        hp['ffn_hidden_size'], hp['num_heads'], hp['enc_layers'],
+        hp['enc_ffn_kernel_size'], hp['dropout'], prenet=hp['enc_prenet'], pre_ln=hp['enc_pre_ln']),
+}
+
+def fill_with_neg_inf2(t):
+    """FP16-compatible function that fills a tensor with -inf."""
+    return t.float().fill_(-1e8).type_as(t)
+
+def expand_states(h, mel2token):
+    h = F.pad(h, [0, 0, 1, 0])
+    mel2token_ = mel2token[..., None].repeat([1, 1, h.shape[-1]])
+    h = torch.gather(h, 1, mel2token_)  # [B, T, H]
+    return h
+
+
+class CodePredictor(nn.Module):
+    def __init__(self, hparams, hidden_size, dec_hidden_size, lm_num_layers, dict_size, code_size):
+        super().__init__()
+        self.hparams = deepcopy(hparams)
+        self.hparams['hidden_size'] = hidden_size
+        self.hidden_size = hidden_size
+        char_dict_size = hparams.get('char_dict_size', 4000)
+        if not hparams.get('lm_use_enc'):
+            self.encoder = nn.Embedding(dict_size, self.hidden_size, padding_idx=0)
+            if hparams.get('mega_use_char', True):
+                self.char_encoder = nn.Embedding(char_dict_size,
+                                                 self.hidden_size, padding_idx=0)
+        else:
+            self.encoder = FS_ENCODERS[self.hparams['encoder_type']](self.hparams, dict_size)
+            if hparams.get('mega_use_char', True):
+                self.char_encoder = FS_ENCODERS[self.hparams['encoder_type']](self.hparams, char_dict_size)
+            if hparams['use_ph_pos_embed']:
+                self.ph_pos_embed = PosEmb(self.hidden_size)
+
+        self.char_empty_embed = nn.Embedding(1, self.hidden_size)
+        if hparams.get('use_bert_input'):
+            self.bert_input_proj = nn.Linear(768, self.hidden_size)
+        self.ling_label_embed_layers = nn.ModuleDict()
+        for k, s in zip(hparams['ling_labels'], hparams['ling_label_dict_size']):
+            self.ling_label_embed_layers[k] = Embedding(s + 3, self.hidden_size, padding_idx=0)
+
+        self.dec_hidden_size = dec_hidden_size
+        self.enc_proj = nn.Linear(self.hidden_size, dec_hidden_size)
+        self.code_emb = Embedding(code_size + 2, dec_hidden_size, 0)
+        self.use_pos_embed = hparams.get('use_pos_embed', False)
+        if self.use_pos_embed:
+            self.embed_positions = SinusoidalPositionalEmbedding(dec_hidden_size, 0, init_size=1024)
+        self.use_post_ln = hparams.get('use_post_ln', False)
+        self.layers = None
+        if not self.use_post_ln:
+            self.layer_norm = LayerNorm(dec_hidden_size)
+        self.code_size = code_size
+        self.project_out_dim = Linear(dec_hidden_size, code_size + 1, bias=True)
+
+    def forward_ling_encoder(
+            self, txt_tokens, ling_feas, char_tokens, ph2char, bert_embed, spk_id, spk_embed, mels_timbre):
+        ph_tokens = txt_tokens
+        hparams = self.hparams
+        ph_nonpadding = (ph_tokens > 0).float()[:, :, None]  # [B, T_phone, 1]
+        x_spk = self.forward_style_embed(spk_embed, spk_id, mels_timbre)
+
+        # enc_ph
+        if not hparams.get('lm_use_enc'):
+            x_ph = self.encoder(ph_tokens)
+            x_ph = x_ph + sum(
+                [self.ling_label_embed_layers[k](ling_feas[k]) for k in hparams['ling_labels']]) \
+                if len(hparams['ling_labels']) > 0 else 0
+            x_ph = x_ph + x_spk
+        else:
+            # enc_ph
+            ph_enc_oembed = sum(
+                [self.ling_label_embed_layers[k](ling_feas[k]) for k in hparams['ling_labels']]) \
+                if len(hparams['ling_labels']) > 0 else 0
+            ph_enc_oembed = ph_enc_oembed + self.ph_pos_embed(
+                torch.arange(0, ph_tokens.shape[1])[None,].to(ph_tokens.device))
+            ph_enc_oembed = ph_enc_oembed + x_spk
+            ph_enc_oembed = ph_enc_oembed * ph_nonpadding
+            x_ph = self.encoder(ph_tokens, other_embeds=ph_enc_oembed)
+
+        # enc_char
+        if char_tokens is not None and ph2char is not None:
+            char_nonpadding = (char_tokens > 0).float()[:, :, None]
+            x_char = self.char_encoder(char_tokens)
+            empty_char = (ph2char > 100000).long()
+            ph2char = ph2char * (1 - empty_char)
+            x_char_phlevel = \
+                expand_states(x_char * char_nonpadding, ph2char) \
+                * (1 - empty_char)[..., None] + \
+                self.char_empty_embed(torch.zeros_like(ph_tokens)) * empty_char[..., None]
+        else:
+            x_char_phlevel = 0
+        # x_ling
+        x_ling = x_ph + x_char_phlevel
+        x_ling = x_ling * ph_nonpadding
+        x_ling = self.enc_proj(x_ling)
+        return x_ling
+
+    def sample_one_step(self, vq_pred):
+        hparams = self.hparams
+        if hparams.get('infer_top_k'):
+            top_k = hparams.get('infer_top_k')
+            temperature = hparams.get('infer_temperature', 1)
+            vq_pred = vq_pred[:, -1] / temperature
+            # optionally crop the logits to only the top k options
+            if top_k is not None:
+                v, _ = torch.topk(vq_pred, min(top_k, vq_pred.size(-1)))
+                vq_pred[vq_pred < v[:, [-1]]] = -float('Inf')
+            # apply softmax to convert logits to (normalized) probabilities
+            probs = F.softmax(vq_pred, dim=-1)
+            # sample from the distribution
+            vq_pred = torch.multinomial(probs, num_samples=1)
+        else:
+            vq_pred = torch.argmax(F.softmax(vq_pred[:, -1], dim=-1), 1)
+        return vq_pred
+
+    def forward_style_embed(self, spk_embed=None, spk_id=None, mel_ref=None):
+        # add spk embed
+        style_embed = 0
+        if self.hparams['use_spk_embed']:
+            style_embed = style_embed + self.spk_embed_proj(spk_embed)[:, None, :]
+        if self.hparams['use_spk_id']:
+            style_embed = style_embed + self.spk_id_proj(spk_id)[:, None, :]
+        if self.hparams['use_spk_enc']:
+            style_embed = style_embed + self.spk_enc(mel_ref)[:, None, :]
+        return style_embed
+
+    def buffered_future_mask(self, tensor):
+        dim = tensor.size(0)
+        if (
+                not hasattr(self, '_future_mask')
+                or self._future_mask is None
+                or self._future_mask.device != tensor.device
+                or self._future_mask.size(0) < dim
+        ):
+            self._future_mask = torch.triu(fill_with_neg_inf2(tensor.new(dim, dim)), 1)
+        return self._future_mask[:dim, :dim]
+
+
+class ARDurPredictor(CodePredictor):
+    def __init__(self, hparams, hidden_size, dec_hidden_size, lm_num_layers, dict_size, code_size, use_rot_embed=True,
+                 op_version=1):
+        super().__init__(hparams, hidden_size, dec_hidden_size, lm_num_layers, dict_size, code_size)
+        self.use_rot_embed = use_rot_embed
+        bias = hparams.get('lm_bias', True)
+        if self.use_rot_embed:
+            self.layers = nn.ModuleList([])
+            self.layers.extend([
+                RotTransformerDecoderLayer(
+                    dec_hidden_size, 0.0, kernel_size=1, ffn_hidden_size=dec_hidden_size * 4,
+                    post_ln=self.use_post_ln, op_version=op_version, bias=bias)
+                for _ in range(lm_num_layers)
+            ])
+        if hparams['dur_model_type'] == 'ar_mse':
+            self.project_out_dim = nn.Sequential(torch.nn.Linear(dec_hidden_size, 1), nn.Softplus())
+        else:
+            self.project_out_dim = torch.nn.Linear(dec_hidden_size, code_size + 1)
+
+    def forward(self, txt_tokens, ling_feas, char_tokens, ph2char, bert_embed,
+                prev_code, spk_id=None, spk_embed=None, mels_timbre=None, mel2ph=None,
+                incremental_state=None, x_ling=None, attn_mask=None, spk_pos_ids_flat=None,
+                prompt_length=None, cache_size=20, streaming=False):
+        x = self.code_emb(prev_code)
+        if x_ling is None:
+            x_ling = self.forward_ling_encoder(
+                txt_tokens, ling_feas, char_tokens, ph2char, bert_embed, spk_id, spk_embed, mels_timbre)
+            x_ling = x_ling.flatten(0, 1)
+            txt_tokens = txt_tokens.flatten(0, 1)
+            x_ling = x_ling[txt_tokens > 0][None]
+
+        # run decoder
+        self_attn_padding_mask = None
+        if self.use_pos_embed:
+            positions = self.embed_positions(
+                prev_code,
+                incremental_state=incremental_state
+            )
+        if incremental_state is not None:
+            x_ling = x_ling[:, x.shape[1] - 1:x.shape[1]]
+            if spk_pos_ids_flat is not None:
+                spk_pos_ids_flat = spk_pos_ids_flat[:, x.shape[1] - 1:x.shape[1]]
+            x = x[:, -1:]
+            if self.use_pos_embed:
+                positions = positions[:, -1:]
+            if streaming:
+                # Shift Pos: query pos is min(cache_size, idx)
+                spk_pos_ids_flat = torch.min(torch.LongTensor([prompt_length + cache_size]).to(x.device),
+                                             spk_pos_ids_flat)
+
+        # # B x T x C -> T x B x C
+        if self.use_pos_embed:
+            x = x + positions
+        x_ling = x_ling[:, :self.hparams['max_tokens']].contiguous()
+        T = min(self.hparams.get('max_tokens_per_item', 1e9), x_ling.shape[1])
+        x_ling = x_ling.reshape(-1, T, x_ling.shape[-1])
+        x = x + x_ling
+        x = x.transpose(0, 1)
+
+        for idx, layer in enumerate(self.layers):
+            if incremental_state is None:
+                self_attn_mask = self.buffered_future_mask(x)
+                if attn_mask is not None:
+                    self_attn_mask = self_attn_mask + (1 - attn_mask.float()) * -1e8
+                self_attn_mask = self_attn_mask.clamp_min(-1e8)
+            else:
+                self_attn_mask = None
+
+            x, attn_weights = layer(
+                x,
+                incremental_state=incremental_state,
+                self_attn_mask=self_attn_mask,
+                self_attn_padding_mask=self_attn_padding_mask,
+                spk_pos_ids_flat=spk_pos_ids_flat
+            )
+
+        if streaming and incremental_state != {}:
+            for k, v in incremental_state.items():
+                if 'attn_state' in k:
+                    prev_key, prev_value = incremental_state[k]['prev_key'], incremental_state[k]['prev_value']
+                    cur_length = prev_key.shape[2]
+                    if cur_length - prompt_length > cache_size:
+                        prev_key = torch.cat((prev_key[:, :, :prompt_length], prev_key[:, :, -cache_size:]), dim=2)
+                        prev_value = torch.cat((prev_value[:, :, :prompt_length], prev_value[:, :, -cache_size:]),
+                                               dim=2)
+                    incremental_state[k]['prev_key'], incremental_state[k]['prev_value'] = prev_key, prev_value
+
+        if not self.use_post_ln:
+            x = self.layer_norm(x)
+        # T x B x C -> B x T x C
+        x = x.transpose(0, 1)
+        x = self.project_out_dim(x)
+        return x
+
+    def infer(self, txt_tokens, ling_feas, char_tokens, ph2char, bert_embed,
+              spk_id=None, spk_embed=None, mels_timbre=None,
+              incremental_state=None, ctx_vqcodes=None, spk_pos_ids_flat=None, return_state=False,
+              first_step_min=0, return_probs=False, first_decoder_inp=None, dur_disturb=0.0, **kwargs):
+        if incremental_state is None:
+            incremental_state = {}
+        x_ling = self.forward_ling_encoder(
+            txt_tokens, ling_feas, char_tokens, ph2char, bert_embed,
+            spk_id, spk_embed, mels_timbre)
+        x_ling = x_ling.flatten(0, 1)
+        txt_tokens_ori = txt_tokens
+        txt_tokens_withpad = txt_tokens = txt_tokens.flatten(0, 1)
+        x_ling = x_ling[txt_tokens > 0][None]
+        txt_tokens = txt_tokens[txt_tokens > 0][None]
+
+        decoded = torch.zeros_like(txt_tokens)
+        decoded = F.pad(decoded, [1, 0], value=self.code_size + 1)
+        if incremental_state != {}:
+            if first_decoder_inp is None:
+                assert ctx_vqcodes is not None
+                decoded[:, :ctx_vqcodes.shape[1]] = ctx_vqcodes
+                ctx_vqcodes = None
+            else:
+                decoded[:, :1] = first_decoder_inp
+        probs = []
+        for step in range(decoded.shape[1] - 1):
+            vq_pred = self(txt_tokens, None, None, None, None,
+                           decoded[:, :step + 1], None, None, None,
+                           incremental_state=incremental_state, x_ling=x_ling,
+                           spk_pos_ids_flat=spk_pos_ids_flat, **kwargs)
+            probs.append(vq_pred.cpu())
+            if ctx_vqcodes is None or step >= ctx_vqcodes.shape[1]:
+                if self.hparams['dur_model_type'] == 'ar_mse':
+                    d = vq_pred[:, -1, 0]
+                    if dur_disturb > 0 and step >= 1:
+                        if random.random() > 0.5:
+                            d = d * (1 + random.random() * dur_disturb)
+                        else:
+                            d = d / (1 + random.random() * dur_disturb)
+                        d = torch.clamp_max(d, self.code_size - 1)
+                    vq_pred = torch.round(d).long()
+                else:
+                    vq_pred = self.sample_one_step(vq_pred)
+                decoded[:, step + 1] = torch.clamp_min(vq_pred, 1)
+                if step == 0:
+                    decoded[:, step + 1] = torch.clamp_min(vq_pred, first_step_min)
+            else:
+                decoded[:, step + 1] = ctx_vqcodes[:, step]
+        decoded = decoded[:, 1:]
+        decoded_2d = torch.zeros_like(txt_tokens_ori)
+        decoded_2d.flatten(0, 1)[txt_tokens_withpad > 0] = decoded
+        if return_state:
+            return decoded_2d, incremental_state
+        if return_probs:
+            return decoded_2d, torch.cat(probs, 1)
+        return decoded_2d
+
+    def streaming_infer(self, txt_tokens, ling_feas, char_tokens, ph2char, bert_embed,
+                        spk_id=None, spk_embed=None, mels_timbre=None,
+                        incremental_state=None, ctx_vqcodes=None, spk_pos_ids_flat=None, return_state=False,
+                        **kwargs):
+        if incremental_state is None:
+            incremental_state = {}
+        x_ling = self.forward_ling_encoder(
+            txt_tokens, ling_feas, char_tokens, ph2char, bert_embed,
+            spk_id, spk_embed, mels_timbre)
+        x_ling = x_ling.flatten(0, 1)
+        txt_tokens_ori = txt_tokens
+        txt_tokens_withpad = txt_tokens = txt_tokens.flatten(0, 1)
+        x_ling = x_ling[txt_tokens > 0][None]
+        txt_tokens = txt_tokens[txt_tokens > 0][None]
+
+        vq_decoded = torch.zeros_like(txt_tokens)
+        vq_decoded = F.pad(vq_decoded, [1, 0], value=self.code_size + 1)
+        if incremental_state != {}:
+            assert ctx_vqcodes is not None
+            vq_decoded[:, :ctx_vqcodes.shape[1]] = ctx_vqcodes
+            ctx_vqcodes = None
+        prompt_length = list(incremental_state.items())[0][1]['prev_key'].shape[2]
+        for step in tqdm(range(vq_decoded.shape[1] - 1), desc='AR Duration Predictor inference...'):
+            vq_pred = self(txt_tokens, None, None, None, None,
+                           vq_decoded[:, :step + 1], None, None, None,
+                           incremental_state=incremental_state, x_ling=x_ling,
+                           spk_pos_ids_flat=spk_pos_ids_flat, prompt_length=prompt_length, streaming=True, **kwargs)
+            if ctx_vqcodes is None or step >= ctx_vqcodes.shape[1]:
+                if self.hparams['dur_model_type'] == 'ar_mse':
+                    vq_pred = torch.round(vq_pred[:, -1, 0]).long()
+                else:
+                    vq_pred = self.sample_one_step(vq_pred)
+                vq_decoded[:, step + 1] = vq_pred
+            else:
+                vq_decoded[:, step + 1] = ctx_vqcodes[:, step]
+        vq_decoded = vq_decoded[:, 1:]
+        vq_decoded_2d = torch.zeros_like(txt_tokens_ori)
+        vq_decoded_2d.flatten(0, 1)[txt_tokens_withpad > 0] = vq_decoded
+        if return_state:
+            return vq_decoded_2d, incremental_state
+        return vq_decoded_2d

xinference/thirdparty/megatts3/tts/modules/ar_dur/commons/layers.py

@@ -0,0 +1,64 @@
+# Copyright 2025 ByteDance and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torch import nn
+
+
+class LayerNorm(torch.nn.LayerNorm):
+    """Layer normalization module.
+    :param int nout: output dim size
+    :param int dim: dimension to be normalized
+    """
+
+    def __init__(self, nout, dim=-1, eps=1e-5):
+        """Construct an LayerNorm object."""
+        super(LayerNorm, self).__init__(nout, eps=eps)
+        self.dim = dim
+
+    def forward(self, x):
+        """Apply layer normalization.
+        :param torch.Tensor x: input tensor
+        :return: layer normalized tensor
+        :rtype torch.Tensor
+        """
+        if self.dim == -1:
+            return super(LayerNorm, self).forward(x)
+        return super(LayerNorm, self).forward(x.transpose(1, -1)).transpose(1, -1)
+
+
+class Reshape(nn.Module):
+    def __init__(self, *args):
+        super(Reshape, self).__init__()
+        self.shape = args
+
+    def forward(self, x):
+        return x.view(self.shape)
+
+
+class Permute(nn.Module):
+    def __init__(self, *args):
+        super(Permute, self).__init__()
+        self.args = args
+
+    def forward(self, x):
+        return x.permute(self.args)
+
+
+def Embedding(num_embeddings, embedding_dim, padding_idx=None):
+    m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx)
+    nn.init.normal_(m.weight, mean=0, std=embedding_dim ** -0.5)
+    if padding_idx is not None:
+        nn.init.constant_(m.weight[padding_idx], 0)
+    return m

xinference/thirdparty/megatts3/tts/modules/ar_dur/commons/nar_tts_modules.py

@@ -0,0 +1,73 @@
+# Copyright 2025 ByteDance and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+
+import torch
+from torch import nn
+
+import torch.nn.functional as F
+
+
+class LengthRegulator(torch.nn.Module):
+    def __init__(self, pad_value=0.0):
+        super(LengthRegulator, self).__init__()
+        self.pad_value = pad_value
+
+    def forward(self, dur, dur_padding=None, alpha=1.0):
+        """
+        Example (no batch dim version):
+            1. dur = [2,2,3]
+            2. token_idx = [[1],[2],[3]], dur_cumsum = [2,4,7], dur_cumsum_prev = [0,2,4]
+            3. token_mask = [[1,1,0,0,0,0,0],
+                             [0,0,1,1,0,0,0],
+                             [0,0,0,0,1,1,1]]
+            4. token_idx * token_mask = [[1,1,0,0,0,0,0],
+                                         [0,0,2,2,0,0,0],
+                                         [0,0,0,0,3,3,3]]
+            5. (token_idx * token_mask).sum(0) = [1,1,2,2,3,3,3]
+
+        :param dur: Batch of durations of each frame (B, T_txt)
+        :param dur_padding: Batch of padding of each frame (B, T_txt)
+        :param alpha: duration rescale coefficient
+        :return:
+            mel2ph (B, T_speech)
+        assert alpha > 0
+        """
+        dur = torch.round(dur.float() * alpha).long()
+        if dur_padding is not None:
+            dur = dur * (1 - dur_padding.long())
+        token_idx = torch.arange(1, dur.shape[1] + 1)[None, :, None].to(dur.device)
+        dur_cumsum = torch.cumsum(dur, 1)
+        dur_cumsum_prev = F.pad(dur_cumsum, [1, -1], mode='constant', value=0)
+
+        pos_idx = torch.arange(dur.sum(-1).max())[None, None].to(dur.device)
+        token_mask = (pos_idx >= dur_cumsum_prev[:, :, None]) & (pos_idx < dur_cumsum[:, :, None])
+        mel2token = (token_idx * token_mask.long()).sum(1)
+        return mel2token
+
+
+class PosEmb(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim) * -emb)
+        self.emb = emb  # TODO
+
+    def forward(self, x):
+        emb = x[:, :, None] * self.emb[None, None, :].to(x.device)
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb