Update Bert-VITS2 Extra

bert_vits2/bert/Erlangshen-MegatronBert-1.3B-Chinese/config.json

@@ -0,0 +1 @@
+{"vocab_size": 21248, "hidden_size": 2048, "num_hidden_layers": 24, "num_attention_heads": 8, "hidden_act": "gelu_new", "intermediate_size": 8192, "hidden_dropout_prob": 0.1, "attention_probs_dropout_prob": 0.1, "max_position_embeddings": 512, "type_vocab_size": 2, "initializer_range": 0.02, "layer_norm_eps": 1e-12, "gradient_checkpointing": false, "position_embedding_type": "absolute", "use_cache": false, "model_type": "megatron-bert"}

bert_vits2/bert_vits2.py

@@ -1,5 +1,6 @@
 import logging
 
+import numpy as np
 import torch
 
 from bert_vits2 import commons
@@ -26,6 +27,7 @@ class Bert_VITS2:
         self.sampling_rate = self.hps_ms.data.sampling_rate
 
         self.bert_model_names = {}
+        self.zh_bert_extra = False
         self.ja_bert_dim = 1024
         self.num_tones = num_tones
 
@@ -35,6 +37,7 @@ class Bert_VITS2:
         self.bert_extra_str_map = {"zh": "", "ja": "", "en": ""}
         self.hps_ms.model.emotion_embedding = None
         if self.version in ["1.0", "1.0.0", "1.0.1"]:
+            self.version = "1.0"
             self.symbols = symbols_legacy
             self.hps_ms.model.n_layers_trans_flow = 3
             self.lang = getattr(self.hps_ms.data, "lang", ["zh"])
@@ -43,6 +46,7 @@ class Bert_VITS2:
             self.text_extra_str_map.update({"zh": "_v100"})
 
         elif self.version in ["1.1.0-transition"]:
+            self.version = "1.1.0-transition"
             self.hps_ms.model.n_layers_trans_flow = 3
             self.lang = getattr(self.hps_ms.data, "lang", ["zh", "ja"])
             self.ja_bert_dim = 768
@@ -52,6 +56,7 @@ class Bert_VITS2:
             self.bert_extra_str_map.update({"ja": "_v111"})
 
         elif self.version in ["1.1", "1.1.0", "1.1.1"]:
+            self.version = "1.1"
             self.hps_ms.model.n_layers_trans_flow = 6
             self.lang = getattr(self.hps_ms.data, "lang", ["zh", "ja"])
             self.ja_bert_dim = 768
@@ -61,6 +66,7 @@ class Bert_VITS2:
             self.bert_extra_str_map.update({"ja": "_v111"})
 
         elif self.version in ["2.0", "2.0.0", "2.0.1", "2.0.2"]:
+            self.version = "2.0"
             self.hps_ms.model.n_layers_trans_flow = 4
             self.lang = getattr(self.hps_ms.data, "lang", ["zh", "ja", "en"])
             self.num_tones = num_tones
@@ -70,6 +76,7 @@ class Bert_VITS2:
             self.bert_extra_str_map.update({"ja": "_v200", "en": "_v200"})
 
         elif self.version in ["2.1", "2.1.0"]:
+            self.version = "2.1"
             self.hps_ms.model.n_layers_trans_flow = 4
             self.hps_ms.model.emotion_embedding = 1
             self.lang = getattr(self.hps_ms.data, "lang", ["zh", "ja", "en"])
@@ -78,6 +85,7 @@ class Bert_VITS2:
             if "en" in self.lang: self.bert_model_names.update({"en": "DEBERTA_V3_LARGE"})
 
         elif self.version in ["2.2", "2.2.0"]:
+            self.version = "2.2"
             self.hps_ms.model.n_layers_trans_flow = 4
             self.hps_ms.model.emotion_embedding = 2
             self.lang = getattr(self.hps_ms.data, "lang", ["zh", "ja", "en"])
@@ -85,20 +93,31 @@ class Bert_VITS2:
             if "ja" in self.lang: self.bert_model_names.update({"ja": "DEBERTA_V2_LARGE_JAPANESE_CHAR_WWM"})
             if "en" in self.lang: self.bert_model_names.update({"en": "DEBERTA_V3_LARGE"})
         elif self.version in ["2.3", "2.3.0"]:
+            self.version = "2.3"
             self.lang = getattr(self.hps_ms.data, "lang", ["zh", "ja", "en"])
             self.num_tones = num_tones
             self.text_extra_str_map.update({"en": "_v230"})
             if "ja" in self.lang: self.bert_model_names.update({"ja": "DEBERTA_V2_LARGE_JAPANESE_CHAR_WWM"})
             if "en" in self.lang: self.bert_model_names.update({"en": "DEBERTA_V3_LARGE"})
+        elif self.version.lower().replace("-", "_") in ["extra", "zh_clap"]:
+            self.version = "extra"
+            self.hps_ms.model.emotion_embedding = 2
+            self.hps_ms.model.n_layers_trans_flow = 6
+            self.lang = ["zh"]
+            self.num_tones = num_tones
+            self.zh_bert_extra = True
+            self.bert_model_names.update({"zh": "Erlangshen-MegatronBert-1.3B-Chinese"})
+            self.bert_extra_str_map.update({"zh": "_extra"})
         else:
             logging.debug("Version information not found. Loaded as the newest version: v2.3.")
+            self.version = "2.3"
             self.lang = getattr(self.hps_ms.data, "lang", ["zh", "ja", "en"])
             self.num_tones = num_tones
             self.text_extra_str_map.update({"en": "_v230"})
             if "ja" in self.lang: self.bert_model_names.update({"ja": "DEBERTA_V2_LARGE_JAPANESE_CHAR_WWM"})
             if "en" in self.lang: self.bert_model_names.update({"en": "DEBERTA_V3_LARGE"})
 
-        if "zh" in self.lang:
+        if "zh" in self.lang and "zh" not in self.bert_model_names.keys():
             self.bert_model_names.update({"zh": "CHINESE_ROBERTA_WWM_EXT_LARGE"})
 
         self._symbol_to_id = {s: i for i, s in enumerate(self.symbols)}
@@ -108,7 +127,7 @@ class Bert_VITS2:
     def load_model(self, model_handler):
         self.model_handler = model_handler
 
-        if self.version in ["2.3", "2.3.0"]:
+        if self.version == "2.3":
             self.net_g = SynthesizerTrn_v230(
                 len(symbols),
                 self.hps_ms.data.filter_length // 2 + 1,
@@ -125,6 +144,7 @@ class Bert_VITS2:
                 symbols=self.symbols,
                 ja_bert_dim=self.ja_bert_dim,
                 num_tones=self.num_tones,
+                zh_bert_extra=self.zh_bert_extra,
                 **self.hps_ms.model).to(self.device)
         _ = self.net_g.eval()
         bert_vits2_utils.load_checkpoint(self.model_path, self.net_g, None, skip_optimizer=True, version=self.version)
@@ -156,7 +176,7 @@ class Bert_VITS2:
         del word2ph
         assert bert.shape[-1] == len(phone), phone
 
-        if language_str == "zh":
+        if language_str == "zh" or self.zh_bert_extra:
             zh_bert = bert
             ja_bert = torch.zeros(self.ja_bert_dim, len(phone))
             en_bert = torch.zeros(1024, len(phone))
@@ -180,7 +200,7 @@ class Bert_VITS2:
         language = torch.LongTensor(language)
         return zh_bert, ja_bert, en_bert, phone, tone, language
 
-    def get_emo_(self, reference_audio, emotion):
+    def _get_emo(self, reference_audio, emotion):
         if reference_audio:
             emo = torch.from_numpy(
                 get_emo(reference_audio, self.model_handler.emotion_model,
@@ -191,6 +211,17 @@ class Bert_VITS2:
 
         return emo
 
+    def _get_clap(self, reference_audio, text_prompt):
+        if isinstance(reference_audio, np.ndarray):
+            emo = get_clap_audio_feature(reference_audio, self.model_handler.clap_model,
+                                         self.model_handler.clap_processor, self.device)
+        else:
+            if text_prompt is None: text_prompt = ""
+            emo = get_clap_text_feature(text_prompt, self.model_handler.clap_model,
+                                        self.model_handler.clap_processor, self.device)
+        emo = torch.squeeze(emo, dim=1).unsqueeze(0)
+        return emo
+
     def _infer(self, id, phones, tones, lang_ids, zh_bert, ja_bert, en_bert, sdp_ratio, noise, noisew, length,
                emo=None):
         with torch.no_grad():
@@ -225,18 +256,11 @@ class Bert_VITS2:
         zh_bert, ja_bert, en_bert, phones, tones, lang_ids = self.get_text(text, lang, self.hps_ms, style_text,
                                                                            style_weigth)
 
+        emo = None
         if self.hps_ms.model.emotion_embedding == 1:
-            emo = self.get_emo_(reference_audio, emotion).to(self.device).unsqueeze(0)
+            emo = self._get_emo(reference_audio, emotion).to(self.device).unsqueeze(0)
         elif self.hps_ms.model.emotion_embedding == 2:
-            if isinstance(text_prompt, str):
-                emo = get_clap_text_feature(text_prompt, self.model_handler.clap_model,
-                                            self.model_handler.clap_processor, self.device)
-            else:
-                emo = get_clap_audio_feature(reference_audio, self.model_handler.clap_model,
-                                             self.model_handler.clap_processor, self.device)
-            emo = torch.squeeze(emo, dim=1).unsqueeze(0)
-        else:
-            emo = None
+            emo = self._get_clap(reference_audio, text_prompt)
 
         return self._infer(id, phones, tones, lang_ids, zh_bert, ja_bert, en_bert, sdp_ratio, noise, noisew, length,
                            emo)
@@ -245,54 +269,49 @@ class Bert_VITS2:
                         text_prompt=None, style_text=None, style_weigth=0.7, **kwargs):
         sentences_list = split_by_language(text, self.lang)
 
+        emo = None
         if self.hps_ms.model.emotion_embedding == 1:
-            emo = self.get_emo_(reference_audio, emotion).to(self.device).unsqueeze(0)
+            emo = self._get_emo(reference_audio, emotion).to(self.device).unsqueeze(0)
         elif self.hps_ms.model.emotion_embedding == 2:
-            if isinstance(text_prompt, str):
-                emo = get_clap_text_feature(text_prompt, self.model_handler.clap_model,
-                                            self.model_handler.clap_processor, self.device)
-            else:
-                emo = get_clap_audio_feature(reference_audio, self.model_handler.clap_model,
-                                             self.model_handler.clap_processor, self.device)
-            emo = torch.squeeze(emo, dim=1).unsqueeze(0)
-        else:
-            emo = None
+            emo = self._get_clap(reference_audio, text_prompt)
 
-        tmp_phones, tmp_tones, tmp_lang_ids, tmp_zh_bert, tmp_ja_bert, tmp_en_bert = [], [], [], [], [], []
+        phones, tones, lang_ids, zh_bert, ja_bert, en_bert = [], [], [], [], [], []
 
         for idx, (_text, lang) in enumerate(sentences_list):
             skip_start = idx != 0
             skip_end = idx != len(sentences_list) - 1
-            zh_bert, ja_bert, en_bert, phones, tones, lang_ids = self.get_text(_text, lang, self.hps_ms, style_text,
-                                                                               style_weigth)
+            _zh_bert, _ja_bert, _en_bert, _phones, _tones, _lang_ids = self.get_text(_text, lang, self.hps_ms,
+                                                                                     style_text, style_weigth)
+
             if skip_start:
-                phones = phones[3:]
-                tones = tones[3:]
-                lang_ids = lang_ids[3:]
-                zh_bert = zh_bert[:, 3:]
-                ja_bert = ja_bert[:, 3:]
-                en_bert = en_bert[:, 3:]
+                _phones = _phones[3:]
+                _tones = _tones[3:]
+                _lang_ids = _lang_ids[3:]
+                _zh_bert = _zh_bert[:, 3:]
+                _ja_bert = _ja_bert[:, 3:]
+                _en_bert = _en_bert[:, 3:]
             if skip_end:
-                phones = phones[:-2]
-                tones = tones[:-2]
-                lang_ids = lang_ids[:-2]
-                zh_bert = zh_bert[:, :-2]
-                ja_bert = ja_bert[:, :-2]
-                en_bert = en_bert[:, :-2]
-
-            tmp_phones.append(phones)
-            tmp_tones.append(tones)
-            tmp_lang_ids.append(lang_ids)
-            tmp_zh_bert.append(zh_bert)
-            tmp_ja_bert.append(ja_bert)
-            tmp_en_bert.append(en_bert)
-
-        zh_bert = torch.concatenate(tmp_zh_bert, dim=1)
-        ja_bert = torch.concatenate(tmp_ja_bert, dim=1)
-        en_bert = torch.concatenate(tmp_en_bert, dim=1)
-        phones = torch.concatenate(tmp_phones, dim=0)
-        tones = torch.concatenate(tmp_tones, dim=0)
-        lang_ids = torch.concatenate(tmp_lang_ids, dim=0)
+                _phones = _phones[:-2]
+                _tones = _tones[:-2]
+                _lang_ids = _lang_ids[:-2]
+                _zh_bert = _zh_bert[:, :-2]
+                _ja_bert = _ja_bert[:, :-2]
+                _en_bert = _en_bert[:, :-2]
+
+            phones.append(_phones)
+            tones.append(_tones)
+            lang_ids.append(_lang_ids)
+            zh_bert.append(_zh_bert)
+            ja_bert.append(_ja_bert)
+            en_bert.append(_en_bert)
+
+        zh_bert = torch.cat(zh_bert, dim=1)
+        ja_bert = torch.cat(ja_bert, dim=1)
+        en_bert = torch.cat(en_bert, dim=1)
+        phones = torch.cat(phones, dim=0)
+        tones = torch.cat(tones, dim=0)
+        lang_ids = torch.cat(lang_ids, dim=0)
+        
         audio = self._infer(id, phones, tones, lang_ids, zh_bert, ja_bert, en_bert, sdp_ratio, noise,
                             noisew, length, emo)
 

bert_vits2/model_handler.py

@@ -13,6 +13,7 @@ from bert_vits2.text.japanese_bert import get_bert_feature as ja_bert
 from bert_vits2.text.japanese_bert_v111 import get_bert_feature as ja_bert_v111
 from bert_vits2.text.japanese_bert_v200 import get_bert_feature as ja_bert_v200
 from bert_vits2.text.english_bert_mock_v200 import get_bert_feature as en_bert_v200
+from bert_vits2.text.chinese_bert_extra import get_bert_feature as zh_bert_extra
 
 
 class ModelHandler:
@@ -53,6 +54,10 @@ class ModelHandler:
             "CLAP_HTSAT_FUSED": [
                 "https://huggingface.co/laion/clap-htsat-fused/resolve/main/pytorch_model.bin?download=true",
                 "https://hf-mirror.com/laion/clap-htsat-fused/resolve/main/pytorch_model.bin?download=true",
+            ],
+            "Erlangshen-MegatronBert-1.3B-Chinese": [
+                "https://huggingface.co/IDEA-CCNL/Erlangshen-UniMC-MegatronBERT-1.3B-Chinese/resolve/main/pytorch_model.bin",
+                "https://hf-mirror.com/IDEA-CCNL/Erlangshen-UniMC-MegatronBERT-1.3B-Chinese/resolve/main/pytorch_model.bin",
             ]
         }
 
@@ -66,6 +71,7 @@ class ModelHandler:
             "DEBERTA_V2_LARGE_JAPANESE_CHAR_WWM": "bf0dab8ad87bd7c22e85ec71e04f2240804fda6d33196157d6b5923af6ea1201",
             "WAV2VEC2_LARGE_ROBUST_12_FT_EMOTION_MSP_DIM": "176d9d1ce29a8bddbab44068b9c1c194c51624c7f1812905e01355da58b18816",
             "CLAP_HTSAT_FUSED": "1ed5d0215d887551ddd0a49ce7311b21429ebdf1e6a129d4e68f743357225253",
+            "Erlangshen-MegatronBert-1.3B-Chinese": "3456bb8f2c7157985688a4cb5cecdb9e229cb1dcf785b01545c611462ffe3579",
         }
         self.model_path = {
             "CHINESE_ROBERTA_WWM_EXT_LARGE": os.path.join(config.ABS_PATH,
@@ -78,11 +84,13 @@ class ModelHandler:
                                                                "bert_vits2/bert/deberta-v2-large-japanese-char-wwm"),
             "WAV2VEC2_LARGE_ROBUST_12_FT_EMOTION_MSP_DIM": os.path.join(config.ABS_PATH,
                                                                         "bert_vits2/emotional/wav2vec2-large-robust-12-ft-emotion-msp-dim"),
-            "CLAP_HTSAT_FUSED": os.path.join(config.ABS_PATH, "bert_vits2/emotional/clap-htsat-fused")
+            "CLAP_HTSAT_FUSED": os.path.join(config.ABS_PATH, "bert_vits2/emotional/clap-htsat-fused"),
+            "Erlangshen-MegatronBert-1.3B-Chinese": os.path.join(config.ABS_PATH,
+                                                                 "bert_vits2/bert/Erlangshen-MegatronBert-1.3B-Chinese"),
         }
 
         self.lang_bert_func_map = {"zh": zh_bert, "en": en_bert, "ja": ja_bert, "ja_v111": ja_bert_v111,
-                                   "ja_v200": ja_bert_v200, "en_v200": en_bert_v200}
+                                   "ja_v200": ja_bert_v200, "en_v200": en_bert_v200, "zh_extra": zh_bert_extra}
 
         self.bert_models = {}  # Value: (tokenizer, model, reference_count)
         self.emotion = None

bert_vits2/models.py

@@ -287,6 +287,7 @@ class TextEncoder(nn.Module):
                  ja_bert_dim=1024,
                  num_tones=None,
                  emotion_embedding=1,
+                 zh_bert_extra=False,
                  ):
         super().__init__()
         self.n_vocab = n_vocab
@@ -305,9 +306,13 @@ class TextEncoder(nn.Module):
         self.language_emb = nn.Embedding(num_languages, hidden_channels)
         nn.init.normal_(self.language_emb.weight, 0.0, hidden_channels ** -0.5)
         self.bert_proj = nn.Conv1d(1024, hidden_channels, 1)
+        self.zh_bert_extra = zh_bert_extra
+        if self.zh_bert_extra:
+            self.bert_pre_proj = nn.Conv1d(2048, 1024, 1)
         self.ja_bert_proj = nn.Conv1d(ja_bert_dim, hidden_channels, 1)
         self.en_bert_proj = nn.Conv1d(1024, hidden_channels, 1)
         self.emotion_embedding = emotion_embedding
+
         if self.emotion_embedding == 1:
             self.emo_proj = nn.Linear(1024, 1024)
             self.emo_quantizer = VectorQuantize(
@@ -356,20 +361,24 @@ class TextEncoder(nn.Module):
         self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
 
     def forward(self, x, x_lengths, tone, language, zh_bert, ja_bert, en_bert, emo=None, sid=None, g=None):
-        zh_bert_emb = self.bert_proj(zh_bert).transpose(1, 2)
-        ja_bert_emb = self.ja_bert_proj(ja_bert).transpose(1, 2)
-        en_bert_emb = self.en_bert_proj(en_bert).transpose(1, 2)
-        x = self.emb(x) + self.tone_emb(tone) + self.language_emb(language) + zh_bert_emb + ja_bert_emb + en_bert_emb
+        x = self.emb(x) + self.tone_emb(tone) + self.language_emb(language)
+
+        if self.zh_bert_extra:
+            zh_bert = self.bert_pre_proj(zh_bert)
+        x += self.bert_proj(zh_bert).transpose(1, 2)
+        x += self.ja_bert_proj(ja_bert).transpose(1, 2)
+        x += self.en_bert_proj(en_bert).transpose(1, 2)
 
+        x *= math.sqrt(self.hidden_channels)  # [b, t, h]
         if self.emotion_embedding == 1:
-            emo = emo.to(zh_bert_emb.device)
+            # emo = emo.to(zh_bert_emb.device)
             if emo.size(-1) == 1024:
                 emo_emb = self.emo_proj(emo.unsqueeze(1))
                 emo_emb_ = []
                 for i in range(emo_emb.size(0)):
                     temp_emo_emb, _, _ = self.emo_quantizer(
-                    emo_emb[i].unsqueeze(0)
-                )
+                        emo_emb[i].unsqueeze(0).to(emo.device)
+                    )
                     emo_emb_.append(temp_emo_emb)
                 emo_emb = torch.cat(emo_emb_, dim=0).to(emo_emb.device)
             else:
@@ -694,6 +703,7 @@ class SynthesizerTrn(nn.Module):
                  ja_bert_dim=1024,
                  num_tones=None,
                  emotion_embedding=False,
+                 zh_bert_extra=False,
                  **kwargs):
 
         super().__init__()
@@ -738,7 +748,8 @@ class SynthesizerTrn(nn.Module):
                                  symbols=symbols,
                                  ja_bert_dim=ja_bert_dim,
                                  num_tones=num_tones,
-                                 emotion_embedding=self.emotion_embedding
+                                 emotion_embedding=self.emotion_embedding,
+                                 zh_bert_extra=zh_bert_extra,
                                  )
         self.dec = Generator(inter_channels, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates,
                              upsample_initial_channel, upsample_kernel_sizes, gin_channels=gin_channels)
@@ -769,8 +780,8 @@ class SynthesizerTrn(nn.Module):
             g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
         x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths, tone, language, zh_bert, ja_bert, en_bert, emo, sid, g=g)
         logw = self.sdp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w) * (sdp_ratio) + self.dp(x, x_mask,
-                                                                                                         g=g) * (
-                       1 - sdp_ratio)
+                                                                                                             g=g) * (
+                               1 - sdp_ratio)
         w = torch.exp(logw) * x_mask * length_scale
         w_ceil = torch.ceil(w)
         y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()

bert_vits2/models_v230.py

@@ -374,19 +374,13 @@ class TextEncoder(nn.Module):
         self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
 
     def forward(self, x, x_lengths, tone, language, zh_bert, ja_bert, en_bert, g=None):
-        zh_bert_emb = self.bert_proj(zh_bert).transpose(1, 2)
-        ja_bert_emb = self.ja_bert_proj(ja_bert).transpose(1, 2)
-        en_bert_emb = self.en_bert_proj(en_bert).transpose(1, 2)
-        x = (
-                    self.emb(x)
-                    + self.tone_emb(tone)
-                    + self.language_emb(language)
-                    + zh_bert_emb
-                    + ja_bert_emb
-                    + en_bert_emb
-            ) * math.sqrt(
-            self.hidden_channels
-        )  # [b, t, h]
+        x = self.emb(x) + self.tone_emb(tone) + self.language_emb(language)
+        
+        x +=self.bert_proj(zh_bert).transpose(1, 2)
+        x += self.ja_bert_proj(ja_bert).transpose(1, 2)
+        x += self.en_bert_proj(en_bert).transpose(1, 2)
+
+        x *= math.sqrt(self.hidden_channels)  # [b, t, h]
         x = torch.transpose(x, 1, -1)  # [b, h, t]
         x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(
             x.dtype
@@ -940,7 +934,7 @@ class SynthesizerTrn(nn.Module):
             sid,
             tone,
             language,
-            bert,
+            zh_bert,
             ja_bert,
             en_bert,
             noise_scale=0.667,
@@ -958,7 +952,7 @@ class SynthesizerTrn(nn.Module):
         else:
             g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
         x, m_p, logs_p, x_mask = self.enc_p(
-            x, x_lengths, tone, language, bert, ja_bert, en_bert, g=g
+            x, x_lengths, tone, language, zh_bert, ja_bert, en_bert, g=g
         )
         logw = self.sdp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w) * (
             sdp_ratio

bert_vits2/text/chinese_bert_extra.py

@@ -0,0 +1,62 @@
+import torch
+
+from utils.config_manager import global_config
+
+
+def get_bert_feature(text, word2ph, tokenizer, model, device=global_config.DEVICE, style_text=None, style_weight=0.7, **kwargs):
+    with torch.no_grad():
+        inputs = tokenizer(text, return_tensors='pt')
+        for i in inputs:
+            inputs[i] = inputs[i].to(device)
+        res = model(**inputs, output_hidden_states=True)
+        res = torch.nn.functional.normalize(
+            torch.cat(res["hidden_states"][-3:-2], -1)[0], dim=0
+        ).cpu()
+        if style_text:
+            style_inputs = tokenizer(style_text, return_tensors="pt")
+            for i in style_inputs:
+                style_inputs[i] = style_inputs[i].to(device)
+            style_res = model(**style_inputs, output_hidden_states=True)
+            style_res = torch.nn.functional.normalize(
+                torch.cat(style_res["hidden_states"][-3:-2], -1)[0], dim=0
+            ).cpu()
+            style_res_mean = style_res.mean(0)
+
+    assert len(word2ph) == len(text) + 2
+    word2phone = word2ph
+    phone_level_feature = []
+    for i in range(len(word2phone)):
+        if style_text:
+            repeat_feature = (
+                    res[i].repeat(word2phone[i], 1) * (1 - style_weight)
+                    + style_res_mean.repeat(word2phone[i], 1) * style_weight
+            )
+        else:
+            repeat_feature = res[i].repeat(word2phone[i], 1)
+        phone_level_feature.append(repeat_feature)
+
+    phone_level_feature = torch.cat(phone_level_feature, dim=0)
+
+    return phone_level_feature.T
+
+
+if __name__ == '__main__':
+
+    word_level_feature = torch.rand(38, 2048) # 12个词,每个词2048维特征
+    word2phone = [1, 2, 1, 2, 2, 1, 2, 2, 1, 2, 2, 1, 2, 2, 2, 2, 2, 1, 1, 2, 2, 1, 2, 2, 2, 2, 1, 2, 2, 2, 2, 2, 1, 2,
+                  2, 2, 2, 1]
+
+    # 计算总帧数
+    total_frames = sum(word2phone)
+    print(word_level_feature.shape)
+    print(word2phone)
+    phone_level_feature = []
+    for i in range(len(word2phone)):
+        print(word_level_feature[i].shape)
+
+        # 对每个词重复word2phone[i]次
+        repeat_feature = word_level_feature[i].repeat(word2phone[i], 1)
+        phone_level_feature.append(repeat_feature)
+
+    phone_level_feature = torch.cat(phone_level_feature, dim=0)
+    print(phone_level_feature.shape) # torch.Size([36, 2048])