Sync files.

AR/data/bucket_sampler.py

@@ -0,0 +1,163 @@
+# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/data/bucket_sampler.py
+# reference: https://github.com/lifeiteng/vall-e
+import itertools
+import math
+import random
+from random import shuffle
+from typing import Iterator
+from typing import Optional
+from typing import TypeVar
+
+import torch
+import torch.distributed as dist
+from torch.utils.data import Dataset
+from torch.utils.data import Sampler
+
+__all__ = [
+    "DistributedBucketSampler",
+]
+
+T_co = TypeVar("T_co", covariant=True)
+
+
+class DistributedBucketSampler(Sampler[T_co]):
+    r"""
+    sort the dataset wrt. input length
+    divide samples into buckets
+    sort within buckets
+    divide buckets into batches
+    sort batches
+    """
+
+    def __init__(
+        self,
+        dataset: Dataset,
+        num_replicas: Optional[int] = None,
+        rank: Optional[int] = None,
+        shuffle: bool = True,
+        seed: int = 0,
+        drop_last: bool = False,
+        batch_size: int = 32,
+    ) -> None:
+        if num_replicas is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            num_replicas = dist.get_world_size() if torch.cuda.is_available() else 1
+        if rank is None:
+            if not dist.is_available():
+                raise RuntimeError("Requires distributed package to be available")
+            rank = dist.get_rank() if torch.cuda.is_available() else 0
+            if torch.cuda.is_available():
+                torch.cuda.set_device(rank)
+        if rank >= num_replicas or rank < 0:
+            raise ValueError(
+                "Invalid rank {}, rank should be in the interval"
+                " [0, {}]".format(rank, num_replicas - 1)
+            )
+        self.dataset = dataset
+        self.num_replicas = num_replicas
+        self.rank = rank
+        self.epoch = 0
+        self.drop_last = drop_last
+        # If the dataset length is evenly divisible by # of replicas, then there
+        # is no need to drop any data, since the dataset will be split equally.
+        if (
+            self.drop_last and len(self.dataset) % self.num_replicas != 0
+        ):  # type: ignore[arg-type]
+            # Split to nearest available length that is evenly divisible.
+            # This is to ensure each rank receives the same amount of data when
+            # using this Sampler.
+            self.num_samples = math.ceil(
+                (len(self.dataset) - self.num_replicas)
+                / self.num_replicas  # type: ignore[arg-type]
+            )
+        else:
+            self.num_samples = math.ceil(
+                len(self.dataset) / self.num_replicas
+            )  # type: ignore[arg-type]
+        self.total_size = self.num_samples * self.num_replicas
+        self.shuffle = shuffle
+        self.seed = seed
+        self.batch_size = batch_size
+        self.id_with_length = self._get_sample_lengths()
+        self.id_buckets = self.make_buckets(bucket_width=2.0)
+
+    def _get_sample_lengths(self):
+        id_with_lengths = []
+        for i in range(len(self.dataset)):
+            id_with_lengths.append((i, self.dataset.get_sample_length(i)))
+        id_with_lengths.sort(key=lambda x: x[1])
+        return id_with_lengths
+
+    def make_buckets(self, bucket_width: float = 2.0):
+        buckets = []
+        cur = []
+        max_sec = bucket_width
+        for id, sec in self.id_with_length:
+            if sec < max_sec:
+                cur.append(id)
+            else:
+                buckets.append(cur)
+                cur = [id]
+                max_sec += bucket_width
+        if len(cur) > 0:
+            buckets.append(cur)
+        return buckets
+
+    def __iter__(self) -> Iterator[T_co]:
+        if self.shuffle:
+            # deterministically shuffle based on epoch and seed
+            g = torch.Generator()
+            g.manual_seed(self.seed + self.epoch)
+            random.seed(self.epoch + self.seed)
+            shuffled_bucket = []
+            for buc in self.id_buckets:
+                buc_copy = buc.copy()
+                shuffle(buc_copy)
+                shuffled_bucket.append(buc_copy)
+            grouped_batch_size = self.batch_size * self.num_replicas
+            shuffled_bucket = list(itertools.chain(*shuffled_bucket))
+            n_batch = int(math.ceil(len(shuffled_bucket) / grouped_batch_size))
+            batches = [
+                shuffled_bucket[b * grouped_batch_size : (b + 1) * grouped_batch_size]
+                for b in range(n_batch)
+            ]
+            shuffle(batches)
+            indices = list(itertools.chain(*batches))
+        else:
+            # type: ignore[arg-type]
+            indices = list(range(len(self.dataset)))
+
+        if not self.drop_last:
+            # add extra samples to make it evenly divisible
+            padding_size = self.total_size - len(indices)
+            if padding_size <= len(indices):
+                indices += indices[:padding_size]
+            else:
+                indices += (indices * math.ceil(padding_size / len(indices)))[
+                    :padding_size
+                ]
+        else:
+            # remove tail of data to make it evenly divisible.
+            indices = indices[: self.total_size]
+        assert len(indices) == self.total_size
+
+        # subsample
+        indices = indices[self.rank : self.total_size : self.num_replicas]
+        assert len(indices) == self.num_samples
+
+        return iter(indices)
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def set_epoch(self, epoch: int) -> None:
+        r"""
+        Sets the epoch for this sampler. When :attr:`shuffle=True`, this ensures all replicas
+        use a different random ordering for each epoch. Otherwise, the next iteration of this
+        sampler will yield the same ordering.
+
+        Args:
+            epoch (int): Epoch number.
+        """
+        self.epoch = epoch

AR/data/data_module.py

@@ -0,0 +1,76 @@
+# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/data/data_module.py
+# reference: https://github.com/lifeiteng/vall-e
+from pytorch_lightning import LightningDataModule
+from moyoyo_tts.AR.data.bucket_sampler import DistributedBucketSampler
+from moyoyo_tts.AR.data.dataset import Text2SemanticDataset
+from torch.utils.data import DataLoader
+
+
+class Text2SemanticDataModule(LightningDataModule):
+    def __init__(
+        self,
+        config,
+        train_semantic_path,
+        train_phoneme_path,
+        dev_semantic_path=None,
+        dev_phoneme_path=None,
+    ):
+        super().__init__()
+        self.config = config
+        self.train_semantic_path = train_semantic_path
+        self.train_phoneme_path = train_phoneme_path
+        self.dev_semantic_path = dev_semantic_path
+        self.dev_phoneme_path = dev_phoneme_path
+        self.num_workers = self.config["data"]["num_workers"]
+
+    def prepare_data(self):
+        pass
+
+    def setup(self, stage=None, output_logs=False):
+        self._train_dataset = Text2SemanticDataset(
+            phoneme_path=self.train_phoneme_path,
+            semantic_path=self.train_semantic_path,
+            max_sec=self.config["data"]["max_sec"],
+            pad_val=self.config["data"]["pad_val"],
+        )
+        self._dev_dataset = self._train_dataset
+        # self._dev_dataset = Text2SemanticDataset(
+        #     phoneme_path=self.dev_phoneme_path,
+        #     semantic_path=self.dev_semantic_path,
+        #     max_sample=self.config['data']['max_eval_sample'],
+        #     max_sec=self.config['data']['max_sec'],
+        #     pad_val=self.config['data']['pad_val'])
+
+    def train_dataloader(self):
+        batch_size=self.config["train"]["batch_size"]//2 if self.config["train"].get("if_dpo",False)==True else self.config["train"]["batch_size"]
+        batch_size = max(min(batch_size,len(self._train_dataset)//4),1)#防止不保存
+        sampler = DistributedBucketSampler(self._train_dataset, batch_size=batch_size)
+        return DataLoader(
+            self._train_dataset,
+            batch_size=batch_size,
+            sampler=sampler,
+            collate_fn=self._train_dataset.collate,
+            num_workers=self.num_workers,
+            persistent_workers=True,
+            prefetch_factor=16,
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self._dev_dataset,
+            batch_size=1,
+            shuffle=False,
+            collate_fn=self._train_dataset.collate,
+            num_workers=max(self.num_workers, 12),
+            persistent_workers=True,
+            prefetch_factor=16,
+        )
+
+    # 这个会使用到嘛？
+    def test_dataloader(self):
+        return DataLoader(
+            self._dev_dataset,
+            batch_size=1,
+            shuffle=False,
+            collate_fn=self._train_dataset.collate,
+        )

AR/data/dataset.py

@@ -0,0 +1,323 @@
+# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/data/dataset.py
+# reference: https://github.com/lifeiteng/vall-e
+import pdb
+import sys
+
+# sys.path.append("/data/docker/liujing04/gpt-vits/mq-vits-s1bert_no_bert")
+import traceback, os
+from typing import Dict
+from typing import List
+
+import numpy as np
+import pandas as pd
+import torch, json
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset
+from transformers import AutoTokenizer
+
+version = os.environ.get('version',None)
+
+from moyoyo_tts.text import cleaned_text_to_sequence
+
+# from config import exp_dir
+
+
+def batch_sequences(sequences: List[np.array], axis: int = 0, pad_value: int = 0):
+    seq = sequences[0]
+    ndim = seq.ndim
+    if axis < 0:
+        axis += ndim
+    dtype = seq.dtype
+    pad_value = dtype.type(pad_value)
+    seq_lengths = [seq.shape[axis] for seq in sequences]
+    max_length = np.max(seq_lengths)
+
+    padded_sequences = []
+    for seq, length in zip(sequences, seq_lengths):
+        padding = (
+            [(0, 0)] * axis + [(0, max_length - length)] + [(0, 0)] * (ndim - axis - 1)
+        )
+        padded_seq = np.pad(seq, padding, mode="constant", constant_values=pad_value)
+        padded_sequences.append(padded_seq)
+    batch = np.stack(padded_sequences)
+    return batch
+
+
+class Text2SemanticDataset(Dataset):
+    """dataset class for text tokens to semantic model training."""
+
+    def __init__(
+        self,
+        phoneme_path: str,
+        semantic_path: str,
+        max_sample: int = None,
+        max_sec: int = 100,
+        pad_val: int = 1024,
+        # min value of phoneme/sec
+        min_ps_ratio: int = 3,
+        # max value of phoneme/sec
+        max_ps_ratio: int = 25,
+    ) -> None:
+        super().__init__()
+
+        self.semantic_data = pd.read_csv(
+            semantic_path, delimiter="\t", encoding="utf-8"
+        )
+        # get dict
+        self.path2 = phoneme_path  # "%s/2-name2text.txt"%exp_dir#phoneme_path
+        self.path3 = "%s/3-bert" % (
+            os.path.dirname(phoneme_path)
+        )  # "%s/3-bert"%exp_dir#bert_dir
+        self.path6 = semantic_path  # "%s/6-name2semantic.tsv"%exp_dir#semantic_path
+        assert os.path.exists(self.path2)
+        assert os.path.exists(self.path6)
+        self.phoneme_data = {}
+        with open(self.path2, "r", encoding="utf8") as f:
+            lines = f.read().strip("\n").split("\n")
+
+        for line in lines:
+            tmp = line.split("\t")
+            if len(tmp) != 4:
+                continue
+            self.phoneme_data[tmp[0]] = [tmp[1], tmp[2], tmp[3]]
+
+        # self.phoneme_data = np.load(phoneme_path, allow_pickle=True).item()
+        # pad for semantic tokens
+        self.PAD: int = pad_val
+        # self.hz = 25
+        # with open("/data/docker/liujing04/gpt-vits/mq-vits-s1bert_no_bert/configs/s2.json", "r") as f:data = f.read()
+        # data=json.loads(data)["model"]["semantic_frame_rate"]#50hz
+        # self.hz=int(data[:-2])#
+        self.hz = int(os.environ.get("hz", "25hz")[:-2])
+
+        # max seconds of semantic token
+        self.max_sec = max_sec
+        self.min_ps_ratio = min_ps_ratio
+        self.max_ps_ratio = max_ps_ratio
+
+        if max_sample is not None:
+            self.semantic_data = self.semantic_data[:max_sample]
+
+        # {idx: (semantic, phoneme)}
+        # semantic list, phoneme list
+        self.semantic_phoneme = []
+        self.item_names = []
+
+        self.inited = False
+
+        if not self.inited:
+            # 调用初始化函数
+            self.init_batch()
+            self.inited = True
+            del self.semantic_data
+            del self.phoneme_data
+        # self.tokenizer = AutoTokenizer.from_pretrained("hfl/chinese-roberta-wwm-ext-large")
+        # self.tokenizer = AutoTokenizer.from_pretrained("/data/docker/liujing04/bert-vits2/Bert-VITS2-master20231106/bert/chinese-roberta-wwm-ext-large")
+
+    def init_batch(self):
+        semantic_data_len = len(self.semantic_data)
+        phoneme_data_len = len(self.phoneme_data.keys())
+        print("semantic_data_len:", semantic_data_len)
+        print("phoneme_data_len:", phoneme_data_len)
+        print(self.semantic_data)
+        idx = 0
+        num_not_in = 0
+        num_deleted_bigger = 0
+        num_deleted_ps = 0
+        for i in range(semantic_data_len):
+            # 先依次遍历
+            # get str
+            item_name = self.semantic_data.iloc[i,0]
+            # print(self.phoneme_data)
+            try:
+                phoneme, word2ph, text = self.phoneme_data[item_name]
+            except Exception:
+                traceback.print_exc()
+                # print(f"{item_name} not in self.phoneme_data !")
+                num_not_in += 1
+                continue
+
+            semantic_str = self.semantic_data.iloc[i,1]
+            # get token list
+            semantic_ids = [int(idx) for idx in semantic_str.split(" ")]
+            # (T), 是否需要��成 (1, T) -> 不需要，因为需要求 len
+            # 过滤掉太长的样本
+            if (
+                len(semantic_ids) > self.max_sec * self.hz
+            ):  #########1###根据token个数推测总时长过滤时长60s（config里）#40*25=1k
+                num_deleted_bigger += 1
+                continue
+            # (T, ), 这个速度不会很慢，所以可以在一开始就处理，无需在 __getitem__ 里面单个处理####
+            phoneme = phoneme.split(" ")
+
+            try:
+                phoneme_ids = cleaned_text_to_sequence(phoneme, version)
+            except:
+                traceback.print_exc()
+                # print(f"{item_name} not in self.phoneme_data !")
+                num_not_in += 1
+                continue
+            # if len(phoneme_ids) >400:###########2：改为恒定限制为semantic/2.5就行
+            if (
+                len(phoneme_ids) > self.max_sec * self.hz / 2.5
+            ):  ###########2：改为恒定限制为semantic/2.5就行
+                num_deleted_ps += 1
+                continue
+            # if len(semantic_ids) > 1000:###########3
+            #     num_deleted_bigger += 1
+            #     continue
+
+            ps_ratio = len(phoneme_ids) / (len(semantic_ids) / self.hz)
+
+            if (
+                ps_ratio > self.max_ps_ratio or ps_ratio < self.min_ps_ratio
+            ):  ##########4#3~25#每秒多少个phone
+                num_deleted_ps += 1
+                # print(item_name)
+                continue
+
+            self.semantic_phoneme.append((semantic_ids, phoneme_ids))
+            idx += 1
+            self.item_names.append(item_name)
+
+        min_num = 100  # 20直接不补#30补了也不存ckpt
+        leng = len(self.semantic_phoneme)
+        if leng < min_num:
+            tmp1 = self.semantic_phoneme
+            tmp2 = self.item_names
+            self.semantic_phoneme = []
+            self.item_names = []
+            for _ in range(max(2, int(min_num / leng))):
+                self.semantic_phoneme += tmp1
+                self.item_names += tmp2
+        if num_not_in > 0:
+            print(f"there are {num_not_in} semantic datas not in phoneme datas")
+        if num_deleted_bigger > 0:
+            print(
+                f"deleted {num_deleted_bigger} audios who's duration are bigger than {self.max_sec} seconds"
+            )
+        if num_deleted_ps > 0:
+            # 4702 for LibriTTS, LirbriTTS 是标注数据, 是否需要筛？=> 需要，有值为 100 的极端值
+            print(
+                f"deleted {num_deleted_ps} audios who's phoneme/sec are bigger than {self.max_ps_ratio} or smaller than {self.min_ps_ratio}"
+            )
+        """
+        there are 31 semantic datas not in phoneme datas
+        deleted 34 audios who's duration are bigger than 54 seconds
+        deleted 3190 audios who's phoneme/sec are bigger than 25 or smaller than 3
+        dataset.__len__(): 366463
+
+        """
+        # 345410 for LibriTTS
+        print("dataset.__len__():", self.__len__())
+
+    def __get_item_names__(self) -> List[str]:
+        return self.item_names
+
+    def __len__(self) -> int:
+        return len(self.semantic_phoneme)
+
+    def __getitem__(self, idx: int) -> Dict:
+        semantic_ids, phoneme_ids = self.semantic_phoneme[idx]
+        item_name = self.item_names[idx]
+        phoneme_ids_len = len(phoneme_ids)
+        # semantic tokens target
+        semantic_ids_len = len(semantic_ids)
+
+        flag = 0
+        path_bert = "%s/%s.pt" % (self.path3, item_name)
+        if os.path.exists(path_bert) == True:
+            bert_feature = torch.load(path_bert, map_location="cpu")
+        else:
+            flag = 1
+        if flag == 1:
+            # bert_feature=torch.zeros_like(phoneme_ids,dtype=torch.float32)
+            bert_feature = None
+        else:
+            assert bert_feature.shape[-1] == len(phoneme_ids)
+        return {
+            "idx": idx,
+            "phoneme_ids": phoneme_ids,
+            "phoneme_ids_len": phoneme_ids_len,
+            "semantic_ids": semantic_ids,
+            "semantic_ids_len": semantic_ids_len,
+            "bert_feature": bert_feature,
+        }
+
+    def get_sample_length(self, idx: int):
+        semantic_ids = self.semantic_phoneme[idx][0]
+        sec = 1.0 * len(semantic_ids) / self.hz
+        return sec
+
+    def collate(self, examples: List[Dict]) -> Dict:
+        sample_index: List[int] = []
+        phoneme_ids: List[torch.Tensor] = []
+        phoneme_ids_lens: List[int] = []
+        semantic_ids: List[torch.Tensor] = []
+        semantic_ids_lens: List[int] = []
+        # return
+
+        for item in examples:
+            sample_index.append(item["idx"])
+            phoneme_ids.append(np.array(item["phoneme_ids"], dtype=np.int64))
+            semantic_ids.append(np.array(item["semantic_ids"], dtype=np.int64))
+            phoneme_ids_lens.append(item["phoneme_ids_len"])
+            semantic_ids_lens.append(item["semantic_ids_len"])
+
+        # pad 0
+        phoneme_ids = batch_sequences(phoneme_ids)
+        semantic_ids = batch_sequences(semantic_ids, pad_value=self.PAD)
+
+        # # convert each batch to torch.tensor
+        phoneme_ids = torch.tensor(phoneme_ids)
+        semantic_ids = torch.tensor(semantic_ids)
+        phoneme_ids_lens = torch.tensor(phoneme_ids_lens)
+        semantic_ids_lens = torch.tensor(semantic_ids_lens)
+        bert_padded = torch.FloatTensor(len(examples), 1024, max(phoneme_ids_lens))
+        bert_padded.zero_()
+
+        for idx, item in enumerate(examples):
+            bert = item["bert_feature"]
+            if bert != None:
+                bert_padded[idx, :, : bert.shape[-1]] = bert
+
+        return {
+            # List[int]
+            "ids": sample_index,
+            # torch.Tensor (B, max_phoneme_length)
+            "phoneme_ids": phoneme_ids,
+            # torch.Tensor (B)
+            "phoneme_ids_len": phoneme_ids_lens,
+            # torch.Tensor (B, max_semantic_ids_length)
+            "semantic_ids": semantic_ids,
+            # torch.Tensor (B)
+            "semantic_ids_len": semantic_ids_lens,
+            # torch.Tensor (B, 1024, max_phoneme_length)
+            "bert_feature": bert_padded,
+        }
+
+
+if __name__ == "__main__":
+    root_dir = "/data/docker/liujing04/gpt-vits/prepare/dump_mix/"
+    dataset = Text2SemanticDataset(
+        phoneme_path=root_dir + "phoneme_train.npy",
+        semantic_path=root_dir + "semantic_train.tsv",
+    )
+
+    batch_size = 12
+    dataloader = DataLoader(
+        dataset, batch_size=batch_size, collate_fn=dataset.collate, shuffle=False
+    )
+    for i, batch in enumerate(dataloader):
+        if i % 1000 == 0:
+            print(i)
+        # if i == 0:
+        #     print('batch["ids"]:', batch["ids"])
+        # print('batch["phoneme_ids"]:', batch["phoneme_ids"],
+        #       batch["phoneme_ids"].shape)
+        # print('batch["phoneme_ids_len"]:', batch["phoneme_ids_len"],
+        #       batch["phoneme_ids_len"].shape)
+        # print('batch["semantic_ids"]:', batch["semantic_ids"],
+        #       batch["semantic_ids"].shape)
+        # print('batch["semantic_ids_len"]:', batch["semantic_ids_len"],
+        #       batch["semantic_ids_len"].shape)

AR/models/t2s_lightning_module.py

@@ -0,0 +1,144 @@
+# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/models/t2s_lightning_module.py
+# reference: https://github.com/lifeiteng/vall-e
+import os
+import sys
+from typing import Dict
+
+import torch
+from pytorch_lightning import LightningModule
+
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+
+from moyoyo_tts.AR.models.t2s_model import Text2SemanticDecoder
+from moyoyo_tts.AR.modules.lr_schedulers import WarmupCosineLRSchedule
+from moyoyo_tts.AR.modules.optim import ScaledAdam
+
+
+class Text2SemanticLightningModule(LightningModule):
+    def __init__(self, config, output_dir, is_train=True):
+        super().__init__()
+        self.config = config
+        self.top_k = 3
+        self.model = Text2SemanticDecoder(config=config, top_k=self.top_k)
+        pretrained_s1 = config.get("pretrained_s1")
+        if pretrained_s1 and is_train:
+            # print(self.load_state_dict(torch.load(pretrained_s1,map_location="cpu")["state_dict"]))
+            print(
+                self.load_state_dict(
+                    torch.load(pretrained_s1, map_location="cpu")["weight"]
+                )
+            )
+        if is_train:
+            self.automatic_optimization = False
+            self.save_hyperparameters()
+            self.eval_dir = output_dir / "eval"
+            self.eval_dir.mkdir(parents=True, exist_ok=True)
+
+    def training_step(self, batch: Dict, batch_idx: int):
+        opt = self.optimizers()
+        scheduler = self.lr_schedulers()
+        forward=self.model.forward if self.config["train"].get("if_dpo",False)==True else self.model.forward_old
+        loss, acc = forward(
+            batch["phoneme_ids"],
+            batch["phoneme_ids_len"],
+            batch["semantic_ids"],
+            batch["semantic_ids_len"],
+            batch["bert_feature"],
+        )
+        self.manual_backward(loss)
+        if batch_idx > 0 and batch_idx % 4 == 0:
+            opt.step()
+            opt.zero_grad()
+            scheduler.step()
+
+        self.log(
+            "total_loss",
+            loss,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            sync_dist=True,
+        )
+        self.log(
+            "lr",
+            scheduler.get_last_lr()[0],
+            on_epoch=True,
+            prog_bar=True,
+            sync_dist=True,
+        )
+        self.log(
+            f"top_{self.top_k}_acc",
+            acc,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            sync_dist=True,
+        )
+
+    def validation_step(self, batch: Dict, batch_idx: int):
+        return
+
+    # # get loss
+    # loss, acc = self.model.forward(
+    #     batch['phoneme_ids'], batch['phoneme_ids_len'],
+    #     batch['semantic_ids'], batch['semantic_ids_len'],
+    #     batch['bert_feature']
+    # )
+    #
+    # self.log(
+    #     "val_total_loss",
+    #     loss,
+    #     on_step=True,
+    #     on_epoch=True,
+    #     prog_bar=True,
+    #     sync_dist=True)
+    # self.log(
+    #     f"val_top_{self.top_k}_acc",
+    #     acc,
+    #     on_step=True,
+    #     on_epoch=True,
+    #     prog_bar=True,
+    #     sync_dist=True)
+    #
+    # # get infer output
+    # semantic_len = batch['semantic_ids'].size(1)
+    # prompt_len = min(int(semantic_len * 0.5), 150)
+    # prompt = batch['semantic_ids'][:, :prompt_len]
+    # pred_semantic = self.model.infer(batch['phoneme_ids'],
+    #                                  batch['phoneme_ids_len'], prompt,
+    #                                  batch['bert_feature']
+    #                                  )
+    # save_name = f'semantic_toks_{batch_idx}.pt'
+    # save_path = os.path.join(self.eval_dir, save_name)
+    # torch.save(pred_semantic.detach().cpu(), save_path)
+
+    def configure_optimizers(self):
+        model_parameters = self.model.parameters()
+        parameters_names = []
+        parameters_names.append(
+            [name_param_pair[0] for name_param_pair in self.model.named_parameters()]
+        )
+        lm_opt = ScaledAdam(
+            model_parameters,
+            lr=0.01,
+            betas=(0.9, 0.95),
+            clipping_scale=2.0,
+            parameters_names=parameters_names,
+            show_dominant_parameters=False,
+            clipping_update_period=1000,
+        )
+
+        return {
+            "optimizer": lm_opt,
+            "lr_scheduler": {
+                "scheduler": WarmupCosineLRSchedule(
+                    lm_opt,
+                    init_lr=self.config["optimizer"]["lr_init"],
+                    peak_lr=self.config["optimizer"]["lr"],
+                    end_lr=self.config["optimizer"]["lr_end"],
+                    warmup_steps=self.config["optimizer"]["warmup_steps"],
+                    total_steps=self.config["optimizer"]["decay_steps"],
+                )
+            },
+        }

AR/models/t2s_lightning_module_onnx.py

@@ -0,0 +1,107 @@
+# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/models/t2s_lightning_module.py
+# reference: https://github.com/lifeiteng/vall-e
+import os, sys
+
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+from typing import Dict
+
+import torch
+from pytorch_lightning import LightningModule
+from moyoyo_tts.AR.models.t2s_model_onnx import Text2SemanticDecoder
+from moyoyo_tts.AR.modules.lr_schedulers import WarmupCosineLRSchedule
+from moyoyo_tts.AR.modules.optim import ScaledAdam
+
+
+class Text2SemanticLightningModule(LightningModule):
+    def __init__(self, config, output_dir, is_train=True):
+        super().__init__()
+        self.config = config
+        self.top_k = 3
+        self.model = Text2SemanticDecoder(config=config, top_k=self.top_k)
+        pretrained_s1 = config.get("pretrained_s1")
+        if pretrained_s1 and is_train:
+            # print(self.load_state_dict(torch.load(pretrained_s1,map_location="cpu")["state_dict"]))
+            print(
+                self.load_state_dict(
+                    torch.load(pretrained_s1, map_location="cpu")["weight"]
+                )
+            )
+        if is_train:
+            self.automatic_optimization = False
+            self.save_hyperparameters()
+            self.eval_dir = output_dir / "eval"
+            self.eval_dir.mkdir(parents=True, exist_ok=True)
+
+    def training_step(self, batch: Dict, batch_idx: int):
+        opt = self.optimizers()
+        scheduler = self.lr_schedulers()
+        loss, acc = self.model.forward(
+            batch["phoneme_ids"],
+            batch["phoneme_ids_len"],
+            batch["semantic_ids"],
+            batch["semantic_ids_len"],
+            batch["bert_feature"],
+        )
+        self.manual_backward(loss)
+        if batch_idx > 0 and batch_idx % 4 == 0:
+            opt.step()
+            opt.zero_grad()
+            scheduler.step()
+
+        self.log(
+            "total_loss",
+            loss,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            sync_dist=True,
+        )
+        self.log(
+            "lr",
+            scheduler.get_last_lr()[0],
+            on_epoch=True,
+            prog_bar=True,
+            sync_dist=True,
+        )
+        self.log(
+            f"top_{self.top_k}_acc",
+            acc,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            sync_dist=True,
+        )
+
+    def validation_step(self, batch: Dict, batch_idx: int):
+        return
+
+    def configure_optimizers(self):
+        model_parameters = self.model.parameters()
+        parameters_names = []
+        parameters_names.append(
+            [name_param_pair[0] for name_param_pair in self.model.named_parameters()]
+        )
+        lm_opt = ScaledAdam(
+            model_parameters,
+            lr=0.01,
+            betas=(0.9, 0.95),
+            clipping_scale=2.0,
+            parameters_names=parameters_names,
+            show_dominant_parameters=False,
+            clipping_update_period=1000,
+        )
+
+        return {
+            "optimizer": lm_opt,
+            "lr_scheduler": {
+                "scheduler": WarmupCosineLRSchedule(
+                    lm_opt,
+                    init_lr=self.config["optimizer"]["lr_init"],
+                    peak_lr=self.config["optimizer"]["lr"],
+                    end_lr=self.config["optimizer"]["lr_end"],
+                    warmup_steps=self.config["optimizer"]["warmup_steps"],
+                    total_steps=self.config["optimizer"]["decay_steps"],
+                )
+            },
+        }

AR/models/t2s_model.py

@@ -0,0 +1,902 @@
+# 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 torch import nn
+from torch.nn import functional as F
+from torchmetrics.classification import MulticlassAccuracy
+from tqdm import tqdm
+
+from moyoyo_tts.AR.models.utils import make_pad_mask
+from moyoyo_tts.AR.models.utils import (
+    topk_sampling,
+    sample,
+    dpo_loss,
+    make_reject_y,
+    get_batch_logps
+)
+from moyoyo_tts.AR.modules.embedding import SinePositionalEmbedding
+from moyoyo_tts.AR.modules.embedding import TokenEmbedding
+from moyoyo_tts.AR.modules.transformer import LayerNorm
+from moyoyo_tts.AR.modules.transformer import TransformerEncoder
+from moyoyo_tts.AR.modules.transformer import TransformerEncoderLayer
+
+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,
+}
+
+# @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)
+
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            attn_bias.masked_fill_(attn_mask, float("-inf"))
+        else:
+            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.masked_fill_(attn_mask, 0)
+        else:
+            attn_mask[attn_mask!=float("-inf")] =0
+            attn_mask[attn_mask==float("-inf")] =1
+            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.permute(2, 0, 1, 3).reshape(batch_size*q_len, self.hidden_dim)
+        attn = attn.view(q_len, batch_size, self.hidden_dim).transpose(1, 0)
+        attn = F.linear(self.to_mask(attn, padding_mask), self.out_w, self.out_b)
+
+        if padding_mask is not None:
+            for i in range(batch_size):
+                # mask = padding_mask[i,:,0]
+                if self.false.device!= padding_mask.device:
+                    self.false = self.false.to(padding_mask.device)
+                idx = torch.where(padding_mask[i,:,0]==self.false)[0]
+                x_item = x[i,idx,:].unsqueeze(0)
+                attn_item = attn[i,idx,:].unsqueeze(0)
+                x_item = x_item + attn_item
+                x_item = F.layer_norm(
+                    x_item, [self.hidden_dim], self.norm_w1, self.norm_b1, self.norm_eps1
+                )
+                x_item = x_item + self.mlp.forward(x_item)
+                x_item = F.layer_norm(
+                    x_item,
+                    [self.hidden_dim],
+                    self.norm_w2,
+                    self.norm_b2,
+                    self.norm_eps2,
+                )
+                x[i,idx,:] = x_item.squeeze(0)
+            x = self.to_mask(x, padding_mask)
+        else:
+            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 decode_next_token(self, x:torch.Tensor, k_cache:torch.Tensor, v_cache:torch.Tensor, attn_mask:Optional[torch.Tensor]=None, torch_sdpa:bool=True):
+        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)
+
+
+        if torch_sdpa:
+            attn = F.scaled_dot_product_attention(q, k, v)
+        else:
+            attn = scaled_dot_product_attention(q, k, v, attn_mask)
+
+        attn = attn.permute(2, 0, 1, 3).reshape(batch_size*q_len, self.hidden_dim)
+        attn = attn.view(q_len, batch_size, self.hidden_dim).transpose(1, 0)
+        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
+
+
+@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: List[torch.Tensor],
+        v_cache: List[torch.Tensor],
+        attn_mask : Optional[torch.Tensor]=None,
+        torch_sdpa:bool=True
+    ):
+        for i in range(self.num_blocks):
+            x, k_cache[i], v_cache[i] = self.blocks[i].decode_next_token(x, k_cache[i], v_cache[i], attn_mask, torch_sdpa)
+        return x, k_cache, v_cache
+
+
+class Text2SemanticDecoder(nn.Module):
+    def __init__(self, config, norm_first=False, top_k=3):
+        super(Text2SemanticDecoder, self).__init__()
+        self.model_dim = config["model"]["hidden_dim"]
+        self.embedding_dim = config["model"]["embedding_dim"]
+        self.num_head = config["model"]["head"]
+        self.num_layers = config["model"]["n_layer"]
+        self.norm_first = norm_first
+        self.vocab_size = config["model"]["vocab_size"]
+        self.phoneme_vocab_size = config["model"]["phoneme_vocab_size"]
+        self.p_dropout = config["model"]["dropout"]
+        self.EOS = config["model"]["EOS"]
+        self.norm_first = norm_first
+        assert self.EOS == self.vocab_size - 1
+        # should be same as num of kmeans bin
+        # assert self.EOS == 1024
+        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.h = TransformerEncoder(
+            TransformerEncoderLayer(
+                d_model=self.model_dim,
+                nhead=self.num_head,
+                dim_feedforward=self.model_dim * 4,
+                dropout=0.1,
+                batch_first=True,
+                norm_first=norm_first,
+            ),
+            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,
+            average="micro",
+            multidim_average="global",
+            ignore_index=self.EOS,
+        )
+
+        blocks = []
+
+        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
+            )
+
+            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 make_input_data(self, x, x_lens, y, y_lens, bert_feature):
+        x = self.ar_text_embedding(x)
+        x = x + self.bert_proj(bert_feature.transpose(1, 2))
+        x = self.ar_text_position(x)
+        x_mask = make_pad_mask(x_lens)
+
+        y_mask = make_pad_mask(y_lens)
+        y_mask_int = y_mask.type(torch.int64)
+        codes = y.type(torch.int64) * (1 - y_mask_int)
+
+        # Training
+        # AR Decoder
+        y, targets = self.pad_y_eos(codes, y_mask_int, eos_id=self.EOS)
+        x_len = x_lens.max()
+        y_len = y_lens.max()
+        y_emb = self.ar_audio_embedding(y)
+        y_pos = self.ar_audio_position(y_emb)
+
+        xy_padding_mask = torch.concat([x_mask, y_mask], dim=1)
+
+        ar_xy_padding_mask = xy_padding_mask
+
+        x_attn_mask = F.pad(
+            torch.zeros((x_len, x_len), dtype=torch.bool, device=x.device),
+            (0, y_len),
+            value=True,
+        )
+        # x_attn_mask[:, x_len]=False
+        y_attn_mask = F.pad(
+            torch.triu(
+                torch.ones(y_len, y_len, dtype=torch.bool, device=x.device),
+                diagonal=1,
+            ),
+            (x_len, 0),
+            value=False,
+        )
+
+        xy_attn_mask = torch.concat([x_attn_mask, y_attn_mask], dim=0)
+        bsz, src_len = x.shape[0], x_len + y_len
+        _xy_padding_mask = (
+            ar_xy_padding_mask.view(bsz, 1, 1, src_len)
+            .expand(-1, self.num_head, -1, -1)
+            .reshape(bsz * self.num_head, 1, src_len)
+        )
+        xy_attn_mask = xy_attn_mask.logical_or(_xy_padding_mask)
+        new_attn_mask = torch.zeros_like(xy_attn_mask, dtype=x.dtype)
+        new_attn_mask.masked_fill_(xy_attn_mask, float("-inf"))
+        xy_attn_mask = new_attn_mask
+        # x 和完整的 y 一次性输入模型
+        xy_pos = torch.concat([x, y_pos], dim=1)
+
+        return xy_pos, xy_attn_mask, targets
+
+    def forward(self, x, x_lens, y, y_lens, bert_feature):
+        """
+        x: phoneme_ids
+        y: semantic_ids
+        """
+
+        reject_y, reject_y_lens = make_reject_y(y, y_lens)
+
+        xy_pos, xy_attn_mask, targets = self.make_input_data(x, x_lens, y, y_lens, bert_feature)
+
+        xy_dec, _ = self.h(
+            (xy_pos, None),
+            mask=xy_attn_mask,
+        )
+        x_len = x_lens.max()
+        logits = self.ar_predict_layer(xy_dec[:, x_len:])
+
+        ###### DPO #############
+        reject_xy_pos, reject_xy_attn_mask, reject_targets = self.make_input_data(x, x_lens, reject_y, reject_y_lens, bert_feature)
+
+        reject_xy_dec, _ = self.h(
+            (reject_xy_pos, None),
+            mask=reject_xy_attn_mask,
+        )
+        x_len = x_lens.max()
+        reject_logits = self.ar_predict_layer(reject_xy_dec[:, x_len:])
+
+        # loss
+        # from feiteng: 每次 duration 越多, 梯度更新也应该更多, 所以用 sum
+
+        loss_1 = F.cross_entropy(logits.permute(0, 2, 1), targets, reduction="sum")
+        acc = self.ar_accuracy_metric(logits.permute(0, 2, 1).detach(), targets).item()
+
+        A_logits, R_logits = get_batch_logps(logits, reject_logits, targets, reject_targets)
+        loss_2, _, _ = dpo_loss(A_logits, R_logits, 0, 0, 0.2, reference_free=True)
+
+        loss = loss_1 + loss_2
+
+        return loss, acc
+
+    def forward_old(self, x, x_lens, y, y_lens, bert_feature):
+        """
+        x: phoneme_ids
+        y: semantic_ids
+        """
+        x = self.ar_text_embedding(x)
+        x = x + self.bert_proj(bert_feature.transpose(1, 2))
+        x = self.ar_text_position(x)
+        x_mask = make_pad_mask(x_lens)
+
+        y_mask = make_pad_mask(y_lens)
+        y_mask_int = y_mask.type(torch.int64)
+        codes = y.type(torch.int64) * (1 - y_mask_int)
+
+        # Training
+        # AR Decoder
+        y, targets = self.pad_y_eos(codes, y_mask_int, eos_id=self.EOS)
+        x_len = x_lens.max()
+        y_len = y_lens.max()
+        y_emb = self.ar_audio_embedding(y)
+        y_pos = self.ar_audio_position(y_emb)
+
+        xy_padding_mask = torch.concat([x_mask, y_mask], dim=1)
+        ar_xy_padding_mask = xy_padding_mask
+
+        x_attn_mask = F.pad(
+            torch.zeros((x_len, x_len), dtype=torch.bool, device=x.device),
+            (0, y_len),
+            value=True,
+        )
+        y_attn_mask = F.pad(
+            torch.triu(
+                torch.ones(y_len, y_len, dtype=torch.bool, device=x.device),
+                diagonal=1,
+            ),
+            (x_len, 0),
+            value=False,
+        )
+        xy_attn_mask = torch.concat([x_attn_mask, y_attn_mask], dim=0)
+        bsz, src_len = x.shape[0], x_len + y_len
+        _xy_padding_mask = (
+            ar_xy_padding_mask.view(bsz, 1, 1, src_len)
+            .expand(-1, self.num_head, -1, -1)
+            .reshape(bsz * self.num_head, 1, src_len)
+        )
+        xy_attn_mask = xy_attn_mask.logical_or(_xy_padding_mask)
+        new_attn_mask = torch.zeros_like(xy_attn_mask, dtype=x.dtype)
+        new_attn_mask.masked_fill_(xy_attn_mask, float("-inf"))
+        xy_attn_mask = new_attn_mask
+        # x 和完整的 y 一次性输入模型
+        xy_pos = torch.concat([x, y_pos], dim=1)
+        xy_dec, _ = self.h(
+            (xy_pos, None),
+            mask=xy_attn_mask,
+        )
+        logits = self.ar_predict_layer(xy_dec[:, x_len:]).permute(0, 2, 1)
+        # loss
+        # from feiteng: 每次 duration 越多, 梯度更新也应该更多, 所以用 sum
+        loss = F.cross_entropy(logits, targets, reduction="sum")
+        acc = self.ar_accuracy_metric(logits.detach(), targets).item()
+        return loss, acc
+
+    # 需要看下这个函数和 forward 的区别以及没有 semantic 的时候 prompts 输入什么
+    def infer(
+            self,
+            x,
+            x_lens,
+            prompts,
+            bert_feature,
+            top_k: int = -100,
+            early_stop_num: int = -1,
+            temperature: float = 1.0,
+    ):
+        x = self.ar_text_embedding(x)
+        x = x + self.bert_proj(bert_feature.transpose(1, 2))
+        x = self.ar_text_position(x)
+
+        # AR Decoder
+        y = prompts
+        prefix_len = y.shape[1]
+        x_len = x.shape[1]
+        x_attn_mask = torch.zeros((x_len, x_len), dtype=torch.bool)
+        stop = False
+        for _ in tqdm(range(1500)):
+            y_emb = self.ar_audio_embedding(y)
+            y_pos = self.ar_audio_position(y_emb)
+            # x 和逐渐增长的 y 一起输入给模型
+            xy_pos = torch.concat([x, y_pos], dim=1)
+            y_len = y.shape[1]
+            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).to(
+                y.device
+            )
+
+            xy_dec, _ = self.h(
+                (xy_pos, None),
+                mask=xy_attn_mask,
+            )
+            logits = self.ar_predict_layer(xy_dec[:, -1])
+            samples = topk_sampling(
+                logits, top_k=top_k, top_p=1.0, temperature=temperature
+            )
+
+            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:
+                # print(torch.argmax(logits, dim=-1)[0] == self.EOS, samples[0, 0] == self.EOS)
+                stop = True
+            if stop:
+                if prompts.shape[1] == y.shape[1]:
+                    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
+            # 本次生成的 semantic_ids 和之前的 y 构成新的 y
+            # print(samples.shape)#[1,1]#第一个1是bs
+            # import os
+            # os._exit(2333)
+            y = torch.concat([y, samples], dim=1)
+        return y
+
+    def pad_y_eos(self, y, y_mask_int, eos_id):
+        targets = F.pad(y, (0, 1), value=0) + eos_id * F.pad(
+            y_mask_int, (0, 1), value=1
+        )
+        # 错位
+        return targets[:, :-1], targets[:, 1:]
+
+    def infer_panel_batch_infer(
+        self,
+        x:List[torch.LongTensor],  #####全部文本token
+        x_lens:torch.LongTensor,
+        prompts:torch.LongTensor,  ####参考音频token
+        bert_feature:List[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,
+    ):
+        if prompts is None:
+            print("Warning: Prompt free is not supported batch_infer! switch to naive_infer")
+            return self.infer_panel_naive_batched(x, x_lens, prompts, bert_feature, top_k=top_k, top_p=top_p, early_stop_num=early_stop_num, temperature=temperature, **kwargs)
+
+
+        max_len = kwargs.get("max_len",x_lens.max())
+        x_list = []
+        for x_item, bert_item in zip(x, bert_feature):
+            # max_len = max(max_len, x_item.shape[0], bert_item.shape[1])
+            x_item = self.ar_text_embedding(x_item.unsqueeze(0))
+            x_item = x_item + self.bert_proj(bert_item.transpose(0, 1).unsqueeze(0))
+            x_item = self.ar_text_position(x_item).squeeze(0)
+            x_item = F.pad(x_item,(0,0,0,max_len-x_item.shape[0]),value=0) if x_item.shape[0]<max_len else x_item
+            x_list.append(x_item)
+        x = torch.stack(x_list, dim=0)
+
+
+        # AR Decoder
+        y = prompts
+
+        x_len = x.shape[1]
+        x_attn_mask = torch.zeros((x_len, x_len), dtype=torch.bool)
+        stop = False
+
+        k_cache = None
+        v_cache = None
+        ###################  first step ##########################
+        if y is not None:
+            y_emb = self.ar_audio_embedding(y)
+            y_len = y_emb.shape[1]
+            prefix_len = y.shape[1]
+            y_lens = torch.LongTensor([y_emb.shape[1]]*y_emb.shape[0]).to(x.device)
+            y_pos = self.ar_audio_position(y_emb)
+            xy_pos = torch.concat([x, y_pos], dim=1)
+            ref_free = False
+        else:
+            y_emb = None
+            y_len = 0
+            prefix_len = 0
+            y_lens = torch.LongTensor([y_len]*x.shape[0]).to(x.device)
+            y_pos = None
+            xy_pos = x
+            y = torch.zeros(x.shape[0], 0, dtype=torch.int, device=x.device)
+            ref_free = True
+
+
+        ##### create mask #####
+        bsz = x.shape[0]
+        src_len = x_len + y_len
+        y_paddind_mask = make_pad_mask(y_lens, y_len)
+        x_paddind_mask = make_pad_mask(x_lens, max_len)
+
+        # (bsz, x_len + y_len)
+        xy_padding_mask = torch.concat([x_paddind_mask, y_paddind_mask], dim=1)
+
+        x_mask = F.pad(
+            x_attn_mask,
+            (0, y_len),  ###xx的纯0扩展到xx纯0+xy纯1，(x,x+y)
+            value=True,
+        )
+        y_mask = F.pad(  ###yy的右上1扩展到左边xy的0,(y,x+y)
+            torch.triu(torch.ones(y_len, y_len, dtype=torch.bool), diagonal=1),
+            (x_len, 0),
+            value=False,
+        )
+
+        xy_mask = torch.concat([x_mask, y_mask], dim=0).view(1 , src_len, src_len).repeat(bsz, 1, 1).to(x.device)
+        _xy_padding_mask = xy_padding_mask.view(bsz, 1, src_len).repeat(1, src_len, 1)
+
+        for i in range(bsz):
+            l = x_lens[i]
+            _xy_padding_mask[i,l:max_len,:]=True
+
+        xy_attn_mask = xy_mask.logical_or(_xy_padding_mask)
+        xy_attn_mask = xy_attn_mask.unsqueeze(1).expand(-1, self.num_head, -1, -1)
+        xy_attn_mask = xy_attn_mask.bool()
+        xy_padding_mask = xy_padding_mask.view(bsz, src_len, 1).expand(-1, -1, self.model_dim)
+
+        ###### decode #####
+        y_list = [None]*y.shape[0]
+        batch_idx_map = list(range(y.shape[0]))
+        idx_list = [None]*y.shape[0]
+        # for idx in tqdm(range(1500)):
+        for idx in range(1500):
+            if idx == 0:
+                xy_dec, k_cache, v_cache = self.t2s_transformer.process_prompt(xy_pos, xy_attn_mask, xy_padding_mask, False)
+            else:
+                xy_dec, k_cache, v_cache = self.t2s_transformer.decode_next_token(xy_pos, k_cache, v_cache, xy_attn_mask, False)
+            logits = self.ar_predict_layer(
+                xy_dec[:, -1]
+            )
+
+            if idx == 0:
+                xy_attn_mask = F.pad(xy_attn_mask[:,:,-1].unsqueeze(-2),(0,1),value=False)
+                logits = logits[:, :-1]
+            else:
+                xy_attn_mask = F.pad(xy_attn_mask,(0,1),value=False)
+
+            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)
+
+            ####### 移除batch中已经生成完毕的序列,进一步优化计算量
+            tokens = torch.argmax(logits, dim=-1)
+            reserved_idx_of_batch_for_y = None
+            if (self.EOS in samples[:, 0]) or \
+                (self.EOS in tokens):  ###如果生成到EOS，则停止
+                    l1 = samples[:, 0]==self.EOS
+                    l2 = tokens==self.EOS
+                    l = l1.logical_or(l2)
+                    removed_idx_of_batch_for_y = torch.where(l==True)[0].tolist()
+                    reserved_idx_of_batch_for_y = torch.where(l==False)[0]
+                    # batch_indexs = torch.tensor(batch_idx_map, device=y.device)[removed_idx_of_batch_for_y]
+                    for i in removed_idx_of_batch_for_y:
+                        batch_index = batch_idx_map[i]
+                        idx_list[batch_index] = idx - 1
+                        y_list[batch_index] = y[i, :-1]
+
+                    batch_idx_map = [batch_idx_map[i] for i in reserved_idx_of_batch_for_y.tolist()]
+
+            # 只保留batch中未生成完毕的序列
+            if reserved_idx_of_batch_for_y is not None:
+                # index = torch.LongTensor(batch_idx_map).to(y.device)
+                y = torch.index_select(y, dim=0, index=reserved_idx_of_batch_for_y)
+                xy_attn_mask = torch.index_select(xy_attn_mask, dim=0, index=reserved_idx_of_batch_for_y)
+                if k_cache is not None :
+                    for i in range(len(k_cache)):
+                        k_cache[i] = torch.index_select(k_cache[i], dim=0, index=reserved_idx_of_batch_for_y)
+                        v_cache[i] = torch.index_select(v_cache[i], dim=0, index=reserved_idx_of_batch_for_y)
+
+
+            if (early_stop_num != -1 and (y.shape[1] - prefix_len) > early_stop_num) or idx==1499:
+                print("use early stop num:", early_stop_num)
+                stop = True
+                for i, batch_index in enumerate(batch_idx_map):
+                    batch_index = batch_idx_map[i]
+                    idx_list[batch_index] = idx
+                    y_list[batch_index] = y[i, :-1]
+
+            if not (None in idx_list):
+                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
+
+            ####################### 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)
+
+        if (None in idx_list):
+            for i in range(x.shape[0]):
+                if idx_list[i] is None:
+                    idx_list[i] = 1500-1  ###如果没有生成到EOS，就用最大长度代替
+
+        if ref_free:
+            return y_list, [0]*x.shape[0]
+        # print(idx_list)
+        return y_list, idx_list
+
+    def infer_panel_naive_batched(self,
+        x:List[torch.LongTensor],  #####全部文本token
+        x_lens:torch.LongTensor,
+        prompts:torch.LongTensor,  ####参考音频token
+        bert_feature:List[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
+        ):
+        y_list = []
+        idx_list = []
+        for i in range(len(x)):
+            y, idx = self.infer_panel_naive(x[i].unsqueeze(0),
+                                                  x_lens[i],
+                                                  prompts[i].unsqueeze(0) if prompts is not None else None,
+                                                  bert_feature[i].unsqueeze(0),
+                                                  top_k,
+                                                  top_p,
+                                                  early_stop_num,
+                                                  temperature,
+                                                  repetition_penalty,
+                                                  **kwargs)
+            y_list.append(y[0])
+            idx_list.append(idx)
+
+        return y_list, idx_list
+
+    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)
+
+        # AR Decoder
+        y = prompts
+
+        x_len = x.shape[1]
+        x_attn_mask = torch.zeros((x_len, x_len), dtype=torch.bool)
+        stop = False
+        # print(1111111,self.num_layers)
+
+        k_cache = None
+        v_cache = None
+        ###################  first step ##########################
+        if y is not 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)
+            ref_free = False
+        else:
+            y_emb = None
+            y_len = 0
+            prefix_len = 0
+            y_pos = None
+            xy_pos = x
+            y = torch.zeros(x.shape[0], 0, dtype=torch.int, device=x.device)
+            ref_free = True
+
+        bsz = x.shape[0]
+        src_len = x_len + y_len
+        x_attn_mask_pad = F.pad(
+            x_attn_mask,
+            (0, y_len),  ###xx的纯0扩展到xx纯0+xy纯1，(x,x+y)
+            value=True,
+        )
+        y_attn_mask = F.pad(  ###yy的右上1扩展到左边xy的0,(y,x+y)
+            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)):
+        for idx in 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
+
+            ####################### 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)
+
+        if ref_free:
+            return y[:, :-1], 0
+        return y[:, :-1], idx - 1
+
+
+    def infer_panel(
+        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
+    ):
+        return self.infer_panel_naive(x, x_lens, prompts, bert_feature, top_k, top_p, early_stop_num, temperature, repetition_penalty, **kwargs)

AR/models/t2s_model_onnx.py

@@ -0,0 +1,338 @@
+# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/models/t2s_model.py
+# reference: https://github.com/lifeiteng/vall-e
+import torch
+from tqdm import tqdm
+
+from moyoyo_tts.AR.modules.embedding_onnx import SinePositionalEmbedding
+from moyoyo_tts.AR.modules.embedding_onnx import TokenEmbedding
+from moyoyo_tts.AR.modules.transformer_onnx import LayerNorm
+from moyoyo_tts.AR.modules.transformer_onnx import TransformerEncoder
+from moyoyo_tts.AR.modules.transformer_onnx import TransformerEncoderLayer
+from torch import nn
+from torch.nn import functional as F
+from torchmetrics.classification import MulticlassAccuracy
+
+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,
+}
+
+inf_tensor_value = torch.FloatTensor([-float("Inf")]).float()
+
+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 = 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 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"))
+
+    logits = logits / max(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)
+
+    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,
+    **sampling_kwargs,
+):
+    probs = logits_to_probs(
+        logits=logits, previous_tokens=previous_tokens, **sampling_kwargs
+    )
+    idx_next = multinomial_sample_one_no_sync(probs)
+    return idx_next, probs
+
+
+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
+
+    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)
+
+
+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):
+        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, 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,
+        }
+
+        y_emb = self.ar_audio_embedding(y)
+
+        cache["y_emb"] = y_emb
+        y_pos = self.ar_audio_position(y_emb)
+
+        xy_pos = torch.concat([x, y_pos], dim=1)
+
+        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:]
+
+        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)
+
+        y = torch.concat([y, samples], dim=1)
+
+        return y, cache["k"], cache["v"], cache["y_emb"], logits, samples
+
+
+class Text2SemanticDecoder(nn.Module):
+    def __init__(self, config, norm_first=False, top_k=3):
+        super(Text2SemanticDecoder, self).__init__()
+        self.model_dim = config["model"]["hidden_dim"]
+        self.embedding_dim = config["model"]["embedding_dim"]
+        self.num_head = config["model"]["head"]
+        self.num_layers = config["model"]["n_layer"]
+        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.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.h = TransformerEncoder(
+            TransformerEncoderLayer(
+                d_model=self.model_dim,
+                nhead=self.num_head,
+                dim_feedforward=self.model_dim * 4,
+                dropout=0.1,
+                batch_first=True,
+                norm_first=norm_first,
+            ),
+            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,
+            average="micro",
+            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)
+
+    def forward(self, x, prompts, bert_feature):
+        early_stop_num = self.early_stop_num
+        prefix_len = prompts.shape[1]
+
+        x = self.onnx_encoder(x, bert_feature)
+        y, k, v, y_emb, stage, x_example = self.first_stage_decoder(x, prompts)
+
+        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
+            if early_stop_num != -1 and (y.shape[1] - prefix_len) > early_stop_num:
+                stop = True
+            if torch.argmax(logits, dim=-1)[0] == self.EOS or samples[0, 0] == self.EOS:
+                stop = True
+            if stop:
+                break
+        y[0, -1] = 0
+        return y, idx
+
+    def infer(self, x, prompts, bert_feature):
+        top_k = self.top_k
+        early_stop_num = self.early_stop_num
+
+        x = self.onnx_encoder(x, bert_feature)
+
+        y = prompts
+        prefix_len = y.shape[1]
+        x_len = x.shape[1]
+        x_example = x[:,:,0] * 0.0
+        x_attn_mask = torch.matmul(x_example.transpose(0, 1), x_example)
+        x_attn_mask = torch.zeros_like(x_attn_mask, dtype=torch.bool)
+
+        stop = False
+        cache = {
+            "all_stage": self.num_layers,
+            "k": [None] * self.num_layers,
+            "v": [None] * self.num_layers,
+            "y_emb": None,
+            "first_infer": 1,
+            "stage": 0,
+        }
+        for idx in range(1500):
+            if cache["first_infer"] == 1:
+                y_emb = self.ar_audio_embedding(y)
+            else:
+                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)
+            if cache["first_infer"] == 1:
+                xy_pos = torch.concat([x, y_pos], dim=1)
+            else:
+                xy_pos = y_pos[:, -1:]
+            y_len = y_pos.shape[1]
+            if cache["first_infer"] == 1:
+                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)
+            else:
+                xy_attn_mask = torch.zeros((1, x_len + y_len), 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=top_k, top_p=1.0, repetition_penalty=1.35)[0].unsqueeze(0)
+            if early_stop_num != -1 and (y.shape[1] - prefix_len) > 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 prompts.shape[1] == y.shape[1]:
+                    y = torch.concat([y, torch.zeros_like(samples)], dim=1)
+                break
+            y = torch.concat([y, samples], dim=1)
+            cache["first_infer"] = 0
+        return y, idx

AR/models/utils.py

@@ -0,0 +1,229 @@
+# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/models/utils.py
+# reference: https://github.com/lifeiteng/vall-e
+import torch
+import torch.nn.functional as F
+from typing import Tuple
+
+def sequence_mask(length, max_length=None):
+    if max_length is None:
+        max_length = length.max()
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
+    """
+    Args:
+      lengths:
+        A 1-D tensor containing sentence lengths.
+      max_len:
+        The length of masks.
+    Returns:
+      Return a 2-D bool tensor, where masked positions
+      are filled with `True` and non-masked positions are
+      filled with `False`.
+
+    #>>> lengths = torch.tensor([1, 3, 2, 5])
+    #>>> make_pad_mask(lengths)
+    tensor([[False,  True,  True,  True,  True],
+            [False, False, False,  True,  True],
+            [False, False,  True,  True,  True],
+            [False, False, False, False, False]])
+    """
+    assert lengths.ndim == 1, lengths.ndim
+    max_len = max(max_len, lengths.max())
+    n = lengths.size(0)
+    seq_range = torch.arange(0, max_len, device=lengths.device)
+    expaned_lengths = seq_range.unsqueeze(0).expand(n, max_len)
+
+    return expaned_lengths >= lengths.unsqueeze(-1)
+
+
+# https://github.com/microsoft/unilm/blob/master/xtune/src/transformers/modeling_utils.py
+def top_k_top_p_filtering(
+    logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1
+):
+    """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+    Args:
+        logits: logits distribution shape (batch size, vocabulary size)
+        if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+        if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+            Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+        Make sure we keep at least min_tokens_to_keep per batch example in the output
+    From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+    """
+    if top_k > 0:
+        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))  # Safety check
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        if min_tokens_to_keep > 1:
+            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        # scatter sorted tensors to original indexing
+        indices_to_remove = sorted_indices_to_remove.scatter(
+            1, sorted_indices, sorted_indices_to_remove
+        )
+        logits[indices_to_remove] = filter_value
+    return logits
+
+
+def topk_sampling(logits, top_k=10, top_p=1.0, temperature=1.0):
+    # temperature: (`optional`) float
+    #     The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+    # top_k: (`optional`) int
+    #     The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50.
+    # top_p: (`optional`) float
+    #     The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.
+
+    # Temperature (higher temperature => more likely to sample low probability tokens)
+    if temperature != 1.0:
+        logits = logits / temperature
+    # Top-p/top-k filtering
+    logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+    # Sample
+    token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
+    return token
+
+
+from typing import Optional, Tuple
+
+
+def multinomial_sample_one_no_sync(
+    probs_sort,
+):  # Does multinomial sampling without a cuda synchronization
+    q = torch.empty_like(probs_sort).exponential_(1)
+    return torch.argmax(probs_sort / q, dim=-1, keepdim=True).to(dtype=torch.int)
+
+
+def logits_to_probs(
+    logits,
+    previous_tokens: Optional[torch.Tensor] = None,
+    temperature: float = 1.0,
+    top_k: Optional[int] = None,
+    top_p: Optional[int] = None,
+    repetition_penalty: float = 1.0,
+):
+    # if previous_tokens is not None:
+    #     previous_tokens = previous_tokens.squeeze()
+    # print(logits.shape,previous_tokens.shape)
+    # pdb.set_trace()
+    if previous_tokens is not None and repetition_penalty != 1.0:
+        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=1, index=sorted_indices, src=sorted_indices_to_remove
+        )
+        logits = logits.masked_fill(indices_to_remove, -float("Inf"))
+
+    logits = logits / max(temperature, 1e-5)
+
+    if top_k is not None:
+        v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+        pivot = v[: , -1].unsqueeze(-1)
+        logits = torch.where(logits < pivot, -float("Inf"), logits)
+
+    probs = torch.nn.functional.softmax(logits, dim=-1)
+    return probs
+
+
+def sample(
+    logits,
+    previous_tokens: Optional[torch.Tensor] = None,
+    **sampling_kwargs,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    probs = logits_to_probs(
+        logits=logits, previous_tokens=previous_tokens, **sampling_kwargs
+    )
+    idx_next = multinomial_sample_one_no_sync(probs)
+    return idx_next, probs
+
+def dpo_loss(policy_chosen_logps: torch.FloatTensor,
+             policy_rejected_logps: torch.FloatTensor,
+             reference_chosen_logps: torch.FloatTensor,
+             reference_rejected_logps: torch.FloatTensor,
+             beta: float,
+             reference_free: bool = False) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
+    pi_logratios = policy_chosen_logps - policy_rejected_logps
+    ref_logratios = reference_chosen_logps - reference_rejected_logps
+
+    if reference_free:
+        ref_logratios = 0
+
+    logits = pi_logratios - ref_logratios
+
+    losses = -F.logsigmoid(beta * logits)
+    chosen_rewards = beta * (policy_chosen_logps - reference_chosen_logps).detach()
+    rejected_rewards = beta * (policy_rejected_logps - reference_rejected_logps).detach()
+
+    return losses.mean(), chosen_rewards, rejected_rewards
+
+def get_batch_logps(logits_target: torch.FloatTensor, logits_reject: torch.FloatTensor, labels_target: torch.LongTensor, labels_reject: torch.LongTensor, average_log_prob: bool = False) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+
+    # dummy token; we'll ignore the losses on these tokens later
+
+    per_token_logps_target = torch.gather(logits_target.log_softmax(-1), dim=2, index=labels_target.unsqueeze(2)).squeeze(2)
+    per_token_logps_reject = torch.gather(logits_reject.log_softmax(-1), dim=2, index=labels_reject.unsqueeze(2)).squeeze(2)
+
+    return per_token_logps_target.sum(-1), per_token_logps_reject.sum(-1)
+
+def make_reject_y(y_o, y_lens):
+    def repeat_P(y):
+        range_idx, _ = torch.randint(0, len(y), size=(2,)).sort()
+        pre = y[:range_idx[0]]
+        shf = y[range_idx[1]:]
+        range_text = y[range_idx[0]:range_idx[1]]
+        new_y = torch.cat([pre, range_text, range_text, shf])
+        return new_y
+    def lost_P(y):
+        range_idx, _ = torch.randint(0, len(y), size=(2,)).sort()
+        pre = y[:range_idx[0]]
+        shf = y[range_idx[1]:]
+        range_text = y[range_idx[0]:range_idx[1]]
+        new_y = torch.cat([pre, shf])
+        return new_y
+    bs = len(y_lens)
+    reject_y = []
+    reject_y_lens = []
+    for b in range(bs):
+        process_item_idx = torch.randint(0, 1, size=(1, ))[0]
+        if process_item_idx == 0:
+            new_y = repeat_P(y_o[b])
+            reject_y.append(new_y)
+            reject_y_lens.append(len(new_y))
+        elif process_item_idx==1:
+            new_y = lost_P(y_o[b])
+            reject_y.append(new_y)
+            reject_y_lens.append(len(new_y))
+    max_length = max(reject_y_lens)
+    for b in range(bs):
+        pad_length = max_length - reject_y_lens[b]
+        reject_y[b] = torch.cat([reject_y[b], torch.zeros(pad_length, dtype=y_o.dtype, device=y_o.device)], dim=0)
+
+    reject_y = torch.stack(reject_y, dim = 0)
+    reject_y_lens = torch.tensor(reject_y_lens, device=y_lens.device)
+
+    return reject_y, reject_y_lens

AR/modules/activation.py

@@ -0,0 +1,429 @@
+# 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 import functional as F
+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 moyoyo_tts.AR.modules.patched_mha_with_cache import multi_head_attention_forward_patched
+
+F.multi_head_attention_forward = multi_head_attention_forward_patched
+
+
+class MultiheadAttention(Module):
+    r"""Allows the model to jointly attend to information
+    from different representation subspaces as described in the paper:
+    `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_.
+
+    Multi-Head Attention is defined as:
+
+    .. math::
+        \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
+
+    where :math:`head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`.
+
+    ``forward()`` will use a special optimized implementation if all of the following
+    conditions are met:
+
+    - self attention is being computed (i.e., ``query``, ``key``, and ``value`` are the same tensor. This
+      restriction will be loosened in the future.)
+    - Either autograd is disabled (using ``torch.inference_mode`` or ``torch.no_grad``) or no tensor argument ``requires_grad``
+    - training is disabled (using ``.eval()``)
+    - dropout is 0
+    - ``add_bias_kv`` is ``False``
+    - ``add_zero_attn`` is ``False``
+    - ``batch_first`` is ``True`` and the input is batched
+    - ``kdim`` and ``vdim`` are equal to ``embed_dim``
+    - at most one of ``key_padding_mask`` or ``attn_mask`` is passed
+    - if a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ is passed, neither ``key_padding_mask``
+      nor ``attn_mask`` is passed
+
+    If the optimized implementation is in use, a
+    `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ can be passed for
+    ``query``/``key``/``value`` to represent padding more efficiently than using a
+    padding mask. In this case, a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_
+    will be returned, and an additional speedup proportional to the fraction of the input
+    that is padding can be expected.
+
+    Args:
+        embed_dim: Total dimension of the model.
+        num_heads: Number of parallel attention heads. Note that ``embed_dim`` will be split
+            across ``num_heads`` (i.e. each head will have dimension ``embed_dim // num_heads``).
+        dropout: Dropout probability on ``attn_output_weights``. Default: ``0.0`` (no dropout).
+        bias: If specified, adds bias to input / output projection layers. Default: ``True``.
+        add_bias_kv: If specified, adds bias to the key and value sequences at dim=0. Default: ``False``.
+        add_zero_attn: If specified, adds a new batch of zeros to the key and value sequences at dim=1.
+            Default: ``False``.
+        kdim: Total number of features for keys. Default: ``None`` (uses ``kdim=embed_dim``).
+        vdim: Total number of features for values. Default: ``None`` (uses ``vdim=embed_dim``).
+        batch_first: If ``True``, then the input and output tensors are provided
+            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
+
+    Examples::
+
+        >>> # xdoctest: +SKIP
+        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
+        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
+
+    """
+    __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]]:
+        r"""
+        Args:
+            query: Query embeddings of shape :math:`(L, E_q)` for unbatched input, :math:`(L, N, E_q)` when ``batch_first=False``
+                or :math:`(N, L, E_q)` when ``batch_first=True``, where :math:`L` is the target sequence length,
+                :math:`N` is the batch size, and :math:`E_q` is the query embedding dimension ``embed_dim``.
+                Queries are compared against key-value pairs to produce the output.
+                See "Attention Is All You Need" for more details.
+            key: Key embeddings of shape :math:`(S, E_k)` for unbatched input, :math:`(S, N, E_k)` when ``batch_first=False``
+                or :math:`(N, S, E_k)` when ``batch_first=True``, where :math:`S` is the source sequence length,
+                :math:`N` is the batch size, and :math:`E_k` is the key embedding dimension ``kdim``.
+                See "Attention Is All You Need" for more details.
+            value: Value embeddings of shape :math:`(S, E_v)` for unbatched input, :math:`(S, N, E_v)` when
+                ``batch_first=False`` or :math:`(N, S, E_v)` when ``batch_first=True``, where :math:`S` is the source
+                sequence length, :math:`N` is the batch size, and :math:`E_v` is the value embedding dimension ``vdim``.
+                See "Attention Is All You Need" for more details.
+            key_padding_mask: If specified, a mask of shape :math:`(N, S)` indicating which elements within ``key``
+                to ignore for the purpose of attention (i.e. treat as "padding"). For unbatched `query`, shape should be :math:`(S)`.
+                Binary and byte masks are supported.
+                For a binary mask, a ``True`` value indicates that the corresponding ``key`` value will be ignored for
+                the purpose of attention. For a float mask, it will be directly added to the corresponding ``key`` value.
+            need_weights: If specified, returns ``attn_output_weights`` in addition to ``attn_outputs``.
+                Default: ``True``.
+            attn_mask: If specified, a 2D or 3D mask preventing attention to certain positions. Must be of shape
+                :math:`(L, S)` or :math:`(N\cdot\text{num\_heads}, L, S)`, where :math:`N` is the batch size,
+                :math:`L` is the target sequence length, and :math:`S` is the source sequence length. A 2D mask will be
+                broadcasted across the batch while a 3D mask allows for a different mask for each entry in the batch.
+                Binary, byte, and float masks are supported. For a binary mask, a ``True`` value indicates that the
+                corresponding position is not allowed to attend. For a byte mask, a non-zero value indicates that the
+                corresponding position is not allowed to attend. For a float mask, the mask values will be added to
+                the attention weight.
+            average_attn_weights: If true, indicates that the returned ``attn_weights`` should be averaged across
+                heads. Otherwise, ``attn_weights`` are provided separately per head. Note that this flag only has an
+                effect when ``need_weights=True``. Default: ``True`` (i.e. average weights across heads)
+
+        Outputs:
+            - **attn_output** - Attention outputs of shape :math:`(L, E)` when input is unbatched,
+              :math:`(L, N, E)` when ``batch_first=False`` or :math:`(N, L, E)` when ``batch_first=True``,
+              where :math:`L` is the target sequence length, :math:`N` is the batch size, and :math:`E` is the
+              embedding dimension ``embed_dim``.
+            - **attn_output_weights** - Only returned when ``need_weights=True``. If ``average_attn_weights=True``,
+              returns attention weights averaged across heads of shape :math:`(L, S)` when input is unbatched or
+              :math:`(N, L, S)`, where :math:`N` is the batch size, :math:`L` is the target sequence length, and
+              :math:`S` is the source sequence length. If ``average_attn_weights=False``, returns attention weights per
+              head of shape :math:`(\text{num\_heads}, L, S)` when input is unbatched or :math:`(N, \text{num\_heads}, L, S)`.
+
+            .. note::
+                `batch_first` argument is ignored for unbatched inputs.
+        """
+        is_batched = query.dim() == 3
+        if key_padding_mask is not None:
+            _kpm_dtype = key_padding_mask.dtype
+            if _kpm_dtype != torch.bool and not torch.is_floating_point(
+                key_padding_mask
+            ):
+                raise AssertionError(
+                    "only bool and floating types of key_padding_mask are supported"
+                )
+        why_not_fast_path = ""
+        if not is_batched:
+            why_not_fast_path = (
+                f"input not batched; expected query.dim() of 3 but got {query.dim()}"
+            )
+        elif query is not key or key is not value:
+            # When lifting this restriction, don't forget to either
+            # enforce that the dtypes all match or test cases where
+            # they don't!
+            why_not_fast_path = "non-self attention was used (query, key, and value are not the same Tensor)"
+        elif self.in_proj_bias is not None and query.dtype != self.in_proj_bias.dtype:
+            why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_bias ({self.in_proj_bias.dtype}) don't match"
+        elif (
+            self.in_proj_weight is not None and query.dtype != self.in_proj_weight.dtype
+        ):
+            # this case will fail anyway, but at least they'll get a useful error message.
+            why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_weight ({self.in_proj_weight.dtype}) don't match"
+        elif self.training:
+            why_not_fast_path = "training is enabled"
+        elif not self.batch_first:
+            why_not_fast_path = "batch_first was not True"
+        elif self.bias_k is not None:
+            why_not_fast_path = "self.bias_k was not None"
+        elif self.bias_v is not None:
+            why_not_fast_path = "self.bias_v was not None"
+        elif self.dropout:
+            why_not_fast_path = f"dropout was {self.dropout}, required zero"
+        elif self.add_zero_attn:
+            why_not_fast_path = "add_zero_attn was enabled"
+        elif not self._qkv_same_embed_dim:
+            why_not_fast_path = "_qkv_same_embed_dim was not True"
+        elif attn_mask is not None:
+            why_not_fast_path = "attn_mask was not None"
+        elif query.is_nested and key_padding_mask is not None:
+            why_not_fast_path = (
+                "key_padding_mask is not supported with NestedTensor input"
+            )
+        elif self.num_heads % 2 == 1:
+            why_not_fast_path = "num_heads is odd"
+        elif torch.is_autocast_enabled():
+            why_not_fast_path = "autocast is enabled"
+
+        if not why_not_fast_path:
+            tensor_args = (
+                query,
+                key,
+                value,
+                self.in_proj_weight,
+                self.in_proj_bias,
+                self.out_proj.weight,
+                self.out_proj.bias,
+            )
+            # We have to use list comprehensions below because TorchScript does not support
+            # generator expressions.
+            if torch.overrides.has_torch_function(tensor_args):
+                why_not_fast_path = "some Tensor argument has_torch_function"
+            elif not all(
+                [
+                    (x is None or x.is_cuda or "cpu" in str(x.device))
+                    for x in tensor_args
+                ]
+            ):
+                why_not_fast_path = "some Tensor argument is neither CUDA nor CPU"
+            elif torch.is_grad_enabled() and any(
+                [x is not None and x.requires_grad for x in tensor_args]
+            ):
+                why_not_fast_path = (
+                    "grad is enabled and at least one of query or the "
+                    "input/output projection weights or biases requires_grad"
+                )
+            if not why_not_fast_path:
+                return torch._native_multi_head_attention(
+                    query,
+                    key,
+                    value,
+                    self.embed_dim,
+                    self.num_heads,
+                    self.in_proj_weight,
+                    self.in_proj_bias,
+                    self.out_proj.weight,
+                    self.out_proj.bias,
+                    key_padding_mask if key_padding_mask is not None else attn_mask,
+                    need_weights,
+                    average_attn_weights,
+                    1
+                    if key_padding_mask is not None
+                    else 0
+                    if attn_mask is not None
+                    else None,
+                )
+
+        any_nested = query.is_nested or key.is_nested or value.is_nested
+        assert not any_nested, (
+            "MultiheadAttention does not support NestedTensor outside of its fast path. "
+            + f"The fast path was not hit because {why_not_fast_path}"
+        )
+
+        if self.batch_first and is_batched:
+            # make sure that the transpose op does not affect the "is" property
+            if key is value:
+                if query is key:
+                    query = key = value = query.transpose(1, 0)
+                else:
+                    query, key = [x.transpose(1, 0) for x in (query, key)]
+                    value = key
+            else:
+                query, key, value = [x.transpose(1, 0) for x in (query, key, value)]
+
+        if not self._qkv_same_embed_dim:
+            attn_output, attn_output_weights = F.multi_head_attention_forward(
+                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,
+                use_separate_proj_weight=True,
+                q_proj_weight=self.q_proj_weight,
+                k_proj_weight=self.k_proj_weight,
+                v_proj_weight=self.v_proj_weight,
+                average_attn_weights=average_attn_weights,
+                cache=cache,
+            )
+        else:
+            attn_output, attn_output_weights = F.multi_head_attention_forward(
+                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,
+            )
+        if self.batch_first and is_batched:
+            return attn_output.transpose(1, 0), attn_output_weights
+        else:
+            return attn_output, attn_output_weights

AR/modules/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 moyoyo_tts.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)

AR/modules/embedding.py

@@ -0,0 +1,81 @@
+# 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.pe = None
+        self.extend_pe(torch.tensor(0.0).expand(1, 4000))
+
+    def extend_pe(self, x):
+        """Reset the positional encodings."""
+        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) -> torch.Tensor:
+        self.extend_pe(x)
+        output = x.unsqueeze(-1) if x.ndim == 2 else x
+        output = output * self.x_scale + self.alpha * self.pe[:, : x.size(1)]
+        return self.dropout(output)

AR/modules/embedding_onnx.py

@@ -0,0 +1,63 @@
+# 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))
+
+    def extend_pe(self, x):
+        position = torch.cumsum(torch.ones_like(x[:,:,0]), dim=1).transpose(0, 1)
+        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)
+        output = x.unsqueeze(-1) if x.ndim == 2 else x
+        output = output * self.x_scale + self.alpha * pe
+        return self.dropout(output)

AR/modules/lr_schedulers.py

@@ -0,0 +1,83 @@
+# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/modules/lr_schedulers.py
+# reference: https://github.com/lifeiteng/vall-e
+import math
+
+import torch
+from torch import nn
+from torch.optim import Adam
+
+
+class WarmupCosineLRSchedule(torch.optim.lr_scheduler._LRScheduler):
+    """
+    Implements Warmup learning rate schedule until 'warmup_steps', going from 'init_lr' to 'peak_lr' for multiple optimizers.
+    """
+
+    def __init__(
+        self,
+        optimizer,
+        init_lr,
+        peak_lr,
+        end_lr,
+        warmup_steps=10000,
+        total_steps=400000,
+        current_step=0,
+    ):
+        self.init_lr = init_lr
+        self.peak_lr = peak_lr
+        self.end_lr = end_lr
+        self.optimizer = optimizer
+        self._warmup_rate = (peak_lr - init_lr) / warmup_steps
+        self._decay_rate = (end_lr - peak_lr) / (total_steps - warmup_steps)
+        self._current_step = current_step
+        self.lr = init_lr
+        self.warmup_steps = warmup_steps
+        self.total_steps = total_steps
+        self._last_lr = [self.lr]
+
+    def set_lr(self, lr):
+        self._last_lr = [g["lr"] for g in self.optimizer.param_groups]
+        for g in self.optimizer.param_groups:
+            # g['lr'] = lr
+            g["lr"] = self.end_lr  ###锁定用线性
+
+    def step(self):
+        if self._current_step < self.warmup_steps:
+            lr = self.init_lr + self._warmup_rate * self._current_step
+
+        elif self._current_step > self.total_steps:
+            lr = self.end_lr
+
+        else:
+            decay_ratio = (self._current_step - self.warmup_steps) / (
+                self.total_steps - self.warmup_steps
+            )
+            if decay_ratio < 0.0 or decay_ratio > 1.0:
+                raise RuntimeError(
+                    "Decay ratio must be in [0.0, 1.0]. Fix LR scheduler settings."
+                )
+            coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
+            lr = self.end_lr + coeff * (self.peak_lr - self.end_lr)
+
+        self.lr = lr = self.end_lr = 0.002  ###锁定用线性###不听话，直接锁定！
+        self.set_lr(lr)
+        self.lr = lr
+        self._current_step += 1
+        return self.lr
+
+
+if __name__ == "__main__":
+    from matplotlib import pyplot as plt
+    m = nn.Linear(10, 10)
+    opt = Adam(m.parameters(), lr=1e-4)
+    s = WarmupCosineLRSchedule(
+        opt, 1e-6, 2e-4, 1e-6, warmup_steps=2000, total_steps=20000, current_step=0
+    )
+    lrs = []
+    for i in range(25000):
+        s.step()
+        lrs.append(s.lr)
+        print(s.lr)
+
+    plt.plot(lrs)
+    plt.plot(range(0, 25000), lrs)
+    plt.show()

AR/modules/optim.py

@@ -0,0 +1,622 @@
+# Copyright      2022  Xiaomi Corp.        (authors: Daniel Povey)
+#
+# See ../LICENSE for clarification regarding multiple authors
+#
+# 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 contextlib
+import logging
+from collections import defaultdict
+from typing import List
+from typing import Tuple
+
+import torch
+from torch import Tensor
+from torch.optim import Optimizer
+
+
+class BatchedOptimizer(Optimizer):
+    """
+    This class adds to class Optimizer the capability to optimize parameters in batches:
+    it will stack the parameters and their grads for you so the optimizer can work
+    on tensors with an extra leading dimension.  This is intended for speed with GPUs,
+    as it reduces the number of kernels launched in the optimizer.
+
+    Args:
+      params:
+    """
+
+    def __init__(self, params, defaults):
+        super(BatchedOptimizer, self).__init__(params, defaults)
+
+    @contextlib.contextmanager
+    def batched_params(self, param_group, group_params_names):
+        """
+        This function returns (technically, yields) a list of
+          of tuples (p, state), where
+        p is a `fake` parameter that is stacked (over axis 0) from real parameters
+        that share the same shape, and its gradient is also stacked;
+        `state` is the state corresponding to this batch of parameters
+        (it will be physically located in the "state" for one of the real
+        parameters, the last one that has any particular shape and dtype).
+
+        This function is decorated as a context manager so that it can
+        write parameters back to their "real" locations.
+
+        The idea is, instead of doing:
+        <code>
+          for p in group["params"]:
+             state = self.state[p]
+             ...
+        </code>
+        you can do:
+        <code>
+          with self.batched_params(group["params"]) as batches:
+             for p, state, p_names in batches:
+                 ...
+        </code>
+
+        Args:
+          group: a parameter group, which is a list of parameters; should be
+                one of self.param_groups.
+          group_params_names: name for each parameter in group,
+                which is List[str].
+        """
+        batches = defaultdict(
+            list
+        )  # `batches` maps from tuple (dtype_as_str,*shape) to list of nn.Parameter
+        batches_names = defaultdict(
+            list
+        )  # `batches` maps from tuple (dtype_as_str,*shape) to list of str
+
+        assert len(param_group) == len(group_params_names)
+        for p, named_p in zip(param_group, group_params_names):
+            key = (str(p.dtype), *p.shape)
+            batches[key].append(p)
+            batches_names[key].append(named_p)
+
+        batches_names_keys = list(batches_names.keys())
+        sorted_idx = sorted(
+            range(len(batches_names)), key=lambda i: batches_names_keys[i])
+        batches_names = [
+            batches_names[batches_names_keys[idx]] for idx in sorted_idx
+        ]
+        batches = [batches[batches_names_keys[idx]] for idx in sorted_idx]
+
+        stacked_params_dict = dict()
+
+        # turn batches into a list, in deterministic order.
+        # tuples will contain tuples of (stacked_param, state, stacked_params_names),
+        # one for each batch in `batches`.
+        tuples = []
+
+        for batch, batch_names in zip(batches, batches_names):
+            p = batch[0]
+            # we arbitrarily store the state in the
+            # state corresponding to the 1st parameter in the
+            # group.  class Optimizer will take care of saving/loading state.
+            state = self.state[p]
+            p_stacked = torch.stack(batch)
+            grad = torch.stack([
+                torch.zeros_like(p) if p.grad is None else p.grad for p in batch
+            ])
+            p_stacked.grad = grad
+            stacked_params_dict[key] = p_stacked
+            tuples.append((p_stacked, state, batch_names))
+
+        yield tuples  # <-- calling code will do the actual optimization here!
+
+        for ((stacked_params, _state, _names), batch) in zip(tuples, batches):
+            for i, p in enumerate(batch):  # batch is list of Parameter
+                p.copy_(stacked_params[i])
+
+
+class ScaledAdam(BatchedOptimizer):
+    """
+     Implements 'Scaled Adam', a variant of Adam where we scale each parameter's update
+     proportional to the norm of that parameter; and also learn the scale of the parameter,
+     in log space, subject to upper and lower limits (as if we had factored each parameter as
+     param = underlying_param * log_scale.exp())
+
+
+     Args:
+          params:  The parameters or param_groups to optimize (like other Optimizer subclasses)
+              lr:  The learning rate.  We will typically use a learning rate schedule that starts
+                   at 0.03 and decreases over time, i.e. much higher than other common
+                   optimizers.
+     clipping_scale: (e.g. 2.0)
+                   A scale for gradient-clipping: if specified, the normalized gradients
+                   over the whole model will be clipped to have 2-norm equal to
+                   `clipping_scale` times the median 2-norm over the most recent period
+                   of `clipping_update_period` minibatches.  By "normalized gradients",
+                   we mean after multiplying by the rms parameter value for this tensor
+                   [for non-scalars]; this is appropriate because our update is scaled
+                   by this quantity.
+            betas: beta1,beta2 are momentum constants for regular momentum, and moving sum-sq grad.
+                   Must satisfy 0 < beta <= beta2 < 1.
+     scalar_lr_scale: A scaling factor on the learning rate, that we use to update the
+                   scale of each parameter tensor and scalar parameters of the mode..
+                   If each parameter were decomposed
+                   as p * p_scale.exp(), where (p**2).mean().sqrt() == 1.0, scalar_lr_scale
+                   would be a the scaling factor on the learning rate of p_scale.
+              eps:  A general-purpose epsilon to prevent division by zero
+    param_min_rms: Minimum root-mean-square value of parameter tensor, for purposes of
+                   learning the scale on the parameters (we'll constrain the rms of each non-scalar
+                   parameter tensor to be >= this value)
+    param_max_rms: Maximum root-mean-square value of parameter tensor, for purposes of
+                   learning the scale on the parameters (we'll constrain the rms of each non-scalar
+                   parameter tensor to be <= this value)
+       scalar_max: Maximum absolute value for scalar parameters (applicable if your
+                   model has any parameters with numel() == 1).
+    size_update_period: The periodicity, in steps, with which we update the size (scale)
+                   of the parameter tensor.  This is provided to save a little time
+                   in the update.
+     clipping_update_period: if clipping_scale is specified, this is the period
+    """
+
+    def __init__(
+            self,
+            params,
+            lr=3e-02,
+            clipping_scale=None,
+            betas=(0.9, 0.98),
+            scalar_lr_scale=0.1,
+            eps=1.0e-08,
+            param_min_rms=1.0e-05,
+            param_max_rms=3.0,
+            scalar_max=10.0,
+            size_update_period=4,
+            clipping_update_period=100,
+            parameters_names=None,
+            show_dominant_parameters=True, ):
+
+        assert parameters_names is not None, (
+            "Please prepare parameters_names,"
+            "which is a List[List[str]]. Each List[str] is for a group"
+            "and each str is for a parameter")
+        defaults = dict(
+            lr=lr,
+            clipping_scale=clipping_scale,
+            betas=betas,
+            scalar_lr_scale=scalar_lr_scale,
+            eps=eps,
+            param_min_rms=param_min_rms,
+            param_max_rms=param_max_rms,
+            scalar_max=scalar_max,
+            size_update_period=size_update_period,
+            clipping_update_period=clipping_update_period, )
+
+        super(ScaledAdam, self).__init__(params, defaults)
+        assert len(self.param_groups) == len(parameters_names)
+        self.parameters_names = parameters_names
+        self.show_dominant_parameters = show_dominant_parameters
+
+    def __setstate__(self, state):
+        super(ScaledAdam, self).__setstate__(state)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        batch = True
+
+        for group, group_params_names in zip(self.param_groups,
+                                             self.parameters_names):
+
+            with self.batched_params(group["params"],
+                                     group_params_names) as batches:
+
+                # batches is list of pairs (stacked_param, state).  stacked_param is like
+                # a regular parameter, and will have a .grad, but the 1st dim corresponds to
+                # a stacking dim, it is not a real dim.
+
+                if (len(batches[0][1]) ==
+                        0):  # if len(first state) == 0: not yet initialized
+                    clipping_scale = 1
+                else:
+                    clipping_scale = self._get_clipping_scale(group, batches)
+
+                for p, state, _ in batches:
+                    # Perform optimization step.
+                    # grad is not going to be None, we handled that when creating the batches.
+                    grad = p.grad
+                    if grad.is_sparse:
+                        raise RuntimeError(
+                            "ScaledAdam optimizer does not support sparse gradients"
+                        )
+                    # State initialization
+                    if len(state) == 0:
+                        self._init_state(group, p, state)
+
+                    self._step_one_batch(group, p, state, clipping_scale)
+
+        return loss
+
+    def _init_state(self, group: dict, p: Tensor, state: dict):
+        """
+        Initializes state dict for parameter 'p'.  Assumes that dim 0 of tensor p
+        is actually the batch dimension, corresponding to batched-together
+        parameters of a given shape.
+
+
+        Args:
+           group:   Dict to look up configuration values.
+               p: The parameter that we are initializing the state for
+           state: Dict from string to whatever state we are initializing
+        """
+        size_update_period = group["size_update_period"]
+
+        state["step"] = 0
+
+        kwargs = {"device": p.device, "dtype": p.dtype}
+
+        # 'delta' implements conventional momentum.  There are
+        # several different kinds of update going on, so rather than
+        # compute "exp_avg" like in Adam, we store and decay a
+        # parameter-change "delta", which combines all forms of
+        # update.  this is equivalent to how it's done in Adam,
+        # except for the first few steps.
+        state["delta"] = torch.zeros_like(
+            p, memory_format=torch.preserve_format)
+
+        batch_size = p.shape[0]
+        numel = p.numel() // batch_size
+        numel = p.numel()
+
+        if numel > 1:
+            # "param_rms" just periodically records the scalar root-mean-square value of
+            # the parameter tensor.
+            # it has a shape like (batch_size, 1, 1, 1, 1)
+            param_rms = (
+                (p**2).mean(dim=list(range(1, p.ndim)), keepdim=True).sqrt())
+            state["param_rms"] = param_rms
+
+            state["scale_exp_avg_sq"] = torch.zeros_like(param_rms)
+            state["scale_grads"] = torch.zeros(size_update_period,
+                                               *param_rms.shape, **kwargs)
+
+        # exp_avg_sq is the weighted sum of scaled gradients. as in Adam.
+        state["exp_avg_sq"] = torch.zeros_like(
+            p, memory_format=torch.preserve_format)
+
+    def _get_clipping_scale(self,
+                            group: dict,
+                            tuples: List[Tuple[Tensor, dict, List[str]]]
+                            ) -> float:
+        """
+        Returns a scalar factor <= 1.0 that dictates gradient clipping, i.e. we will scale the gradients
+        by this amount before applying the rest of the update.
+
+        Args:
+           group: the parameter group, an item in self.param_groups
+           tuples: a list of tuples of (param, state, param_names)
+                where param is a batched set of parameters,
+                with a .grad (1st dim is batch dim)
+                and state is the state-dict where optimization parameters are kept.
+                param_names is a List[str] while each str is name for a parameter
+                in batched set of parameters "param".
+        """
+        assert len(tuples) >= 1
+        clipping_scale = group["clipping_scale"]
+        (first_p, first_state, _) = tuples[0]
+        step = first_state["step"]
+        if clipping_scale is None or step == 0:
+            # no clipping.  return early on step == 0 because the other
+            # parameters' state won't have been initialized yet.
+            return 1.0
+        clipping_update_period = group["clipping_update_period"]
+
+        tot_sumsq = torch.tensor(0.0, device=first_p.device)
+        for (p, state, param_names) in tuples:
+            grad = p.grad
+            if grad.is_sparse:
+                raise RuntimeError(
+                    "ScaledAdam optimizer does not support sparse gradients")
+            if p.numel() == p.shape[0]:  # a batch of scalars
+                tot_sumsq += (grad**2).sum()  # sum() to change shape [1] to []
+            else:
+                tot_sumsq += ((grad * state["param_rms"])**2).sum()
+
+        tot_norm = tot_sumsq.sqrt()
+        if "model_norms" not in first_state:
+            first_state["model_norms"] = torch.zeros(
+                clipping_update_period, device=p.device)
+        first_state["model_norms"][step % clipping_update_period] = tot_norm
+
+        if step % clipping_update_period == 0:
+            # Print some stats.
+            # We don't reach here if step == 0 because we would have returned
+            # above.
+            sorted_norms = first_state["model_norms"].sort()[0].to("cpu")
+            quartiles = []
+            for n in range(0, 5):
+                index = min(
+                    clipping_update_period - 1,
+                    (clipping_update_period // 4) * n, )
+                quartiles.append(sorted_norms[index].item())
+
+            median = quartiles[2]
+            threshold = clipping_scale * median
+            first_state["model_norm_threshold"] = threshold
+            percent_clipped = (first_state["num_clipped"] * 100.0 /
+                               clipping_update_period
+                               if "num_clipped" in first_state else 0.0)
+            first_state["num_clipped"] = 0
+            quartiles = " ".join(["%.3e" % x for x in quartiles])
+            logging.info(
+                f"Clipping_scale={clipping_scale}, grad-norm quartiles {quartiles}, "
+                f"threshold={threshold:.3e}, percent-clipped={percent_clipped:.1f}"
+            )
+
+        if step < clipping_update_period:
+            return 1.0  # We have not yet estimated a norm to clip to.
+        else:
+            try:
+                model_norm_threshold = first_state["model_norm_threshold"]
+            except KeyError:
+                logging.info(
+                    "Warning: model_norm_threshold not in state: possibly "
+                    "you changed config when restarting, adding clipping_scale option?"
+                )
+                return 1.0
+            ans = min(1.0, (model_norm_threshold / (tot_norm + 1.0e-20)).item())
+            if ans < 1.0:
+                first_state["num_clipped"] += 1
+            if ans < 0.1:
+                logging.warn(
+                    f"Scaling gradients by {ans}, model_norm_threshold={model_norm_threshold}"
+                )
+                if self.show_dominant_parameters:
+                    assert p.shape[0] == len(param_names)
+                    self._show_gradient_dominating_parameter(tuples, tot_sumsq)
+            return ans
+
+    def _show_gradient_dominating_parameter(
+            self, tuples: List[Tuple[Tensor, dict, List[str]]],
+            tot_sumsq: Tensor):
+        """
+        Show information of parameter wihch dominanting tot_sumsq.
+
+        Args:
+           tuples: a list of tuples of (param, state, param_names)
+                where param is a batched set of parameters,
+                with a .grad (1st dim is batch dim)
+                and state is the state-dict where optimization parameters are kept.
+                param_names is a List[str] while each str is name for a parameter
+                in batched set of parameters "param".
+            tot_sumsq: sumsq of all parameters. Though it's could be calculated
+                from tuples, we still pass it to save some time.
+        """
+        all_sumsq_orig = {}
+        for (p, state, batch_param_names) in tuples:
+            # p is a stacked batch parameters.
+            batch_grad = p.grad
+            if p.numel() == p.shape[0]:  # a batch of scalars
+                batch_sumsq_orig = batch_grad**2
+                # Dummpy values used by following `zip` statement.
+                batch_rms_orig = torch.ones(p.shape[0])
+            else:
+                batch_rms_orig = state["param_rms"]
+                batch_sumsq_orig = ((batch_grad * batch_rms_orig)**2).sum(
+                    dim=list(range(1, batch_grad.ndim)))
+
+            for name, sumsq_orig, rms, grad in zip(batch_param_names,
+                                                   batch_sumsq_orig,
+                                                   batch_rms_orig, batch_grad):
+
+                proportion_orig = sumsq_orig / tot_sumsq
+                all_sumsq_orig[name] = (proportion_orig, sumsq_orig, rms, grad)
+
+        assert torch.isclose(
+            sum([value[0] for value in all_sumsq_orig.values()]).cpu(),
+            torch.tensor(1.0), )
+        sorted_by_proportion = {
+            k: v
+            for k, v in sorted(
+                all_sumsq_orig.items(),
+                key=lambda item: item[1][0],
+                reverse=True, )
+        }
+        dominant_param_name = next(iter(sorted_by_proportion))
+        (dominant_proportion, dominant_sumsq, dominant_rms,
+         dominant_grad, ) = sorted_by_proportion[dominant_param_name]
+        logging.info(f"Parameter Dominanting tot_sumsq {dominant_param_name}"
+                     f" with proportion {dominant_proportion:.2f},"
+                     f" where dominant_sumsq=(grad_sumsq*orig_rms_sq)"
+                     f"={dominant_sumsq:.3e},"
+                     f" grad_sumsq = {(dominant_grad**2).sum():.3e},"
+                     f" orig_rms_sq={(dominant_rms**2).item():.3e}")
+
+    def _step_one_batch(self,
+                        group: dict,
+                        p: Tensor,
+                        state: dict,
+                        clipping_scale: float):
+        """
+        Do the step for one parameter, which is actually going to be a batch of
+        `real` parameters, with dim 0 as the batch dim.
+        Args:
+                  group:  dict to look up configuration values
+                    p: parameter to update (actually multiple parameters stacked together
+                       as a batch)
+                  state: state-dict for p, to look up the optimizer state
+        """
+        lr = group["lr"]
+        size_update_period = group["size_update_period"]
+        beta1 = group["betas"][0]
+
+        grad = p.grad
+        if clipping_scale != 1.0:
+            grad = grad * clipping_scale
+        step = state["step"]
+        delta = state["delta"]
+
+        delta.mul_(beta1)
+        batch_size = p.shape[0]
+        numel = p.numel() // batch_size
+        if numel > 1:
+            # Update the size/scale of p, and set param_rms
+            scale_grads = state["scale_grads"]
+            scale_grads[step % size_update_period] = (p * grad).sum(
+                dim=list(range(1, p.ndim)), keepdim=True)
+            if step % size_update_period == size_update_period - 1:
+                param_rms = state["param_rms"]  # shape: (batch_size, 1, 1, ..)
+                param_rms.copy_((p**2)
+                                .mean(dim=list(range(1, p.ndim)), keepdim=True)
+                                .sqrt())
+                if step > 0:
+                    # self._size_update() learns the overall scale on the
+                    # parameter, by shrinking or expanding it.
+                    self._size_update(group, scale_grads, p, state)
+
+        if numel == 1:
+            # For parameters with 1 element we just use regular Adam.
+            # Updates delta.
+            self._step_scalar(group, p, state)
+        else:
+            self._step(group, p, state)
+
+        state["step"] = step + 1
+
+    def _size_update(self,
+                     group: dict,
+                     scale_grads: Tensor,
+                     p: Tensor,
+                     state: dict) -> None:
+        """
+               Called only where p.numel() > 1, this updates the scale of the parameter.
+               If we imagine: p =  underlying_param * scale.exp(), and we are doing
+               gradient descent on underlying param and on scale, this function does the update
+               on `scale`.
+
+               Args:
+              group: dict to look up configuration values
+        scale_grads: a tensor of shape (size_update_period, batch_size, 1, 1,...) containing
+                      grads w.r.t. the scales.
+                  p:  The parameter to update
+               state: The state-dict of p
+        """
+
+        param_rms = state["param_rms"]
+        beta1, beta2 = group["betas"]
+        size_lr = group["lr"] * group["scalar_lr_scale"]
+        param_min_rms = group["param_min_rms"]
+        param_max_rms = group["param_max_rms"]
+        eps = group["eps"]
+        step = state["step"]
+        batch_size = p.shape[0]
+
+        size_update_period = scale_grads.shape[0]
+        # correct beta2 for the size update period: we will have
+        # faster decay at this level.
+        beta2_corr = beta2**size_update_period
+
+        scale_exp_avg_sq = state[
+            "scale_exp_avg_sq"]  # shape: (batch_size, 1, 1, ..)
+        scale_exp_avg_sq.mul_(beta2_corr).add_(
+            (scale_grads**2).mean(dim=0),  # mean over dim `size_update_period`
+            alpha=1 - beta2_corr, )  # shape is (batch_size, 1, 1, ...)
+
+        # The 1st time we reach here is when size_step == 1.
+        size_step = (step + 1) // size_update_period
+        bias_correction2 = 1 - beta2_corr**size_step
+        # we don't bother with bias_correction1; this will help prevent divergence
+        # at the start of training.
+
+        denom = scale_exp_avg_sq.sqrt() + eps
+
+        scale_step = (-size_lr * (bias_correction2**0.5) *
+                      scale_grads.sum(dim=0) / denom)
+
+        is_too_small = param_rms < param_min_rms
+        is_too_large = param_rms > param_max_rms
+
+        # when the param gets too small, just don't shrink it any further.
+        scale_step.masked_fill_(is_too_small, 0.0)
+        # when it gets too large, stop it from getting any larger.
+        scale_step.masked_fill_(is_too_large, -size_lr * size_update_period)
+        delta = state["delta"]
+        # the factor of (1-beta1) relates to momentum.
+        delta.add_(p * scale_step, alpha=(1 - beta1))
+
+    def _step(self, group: dict, p: Tensor, state: dict):
+        """
+        This function does the core update of self.step(), in the case where the members of
+        the batch have more than 1 element.
+
+        Args:
+            group: A dict which will be used to look up configuration values
+                p: The parameter to be updated
+             grad: The grad of p
+            state: The state-dict corresponding to parameter p
+
+        This function modifies p.
+        """
+        grad = p.grad
+        lr = group["lr"]
+        beta1, beta2 = group["betas"]
+        eps = group["eps"]
+        param_min_rms = group["param_min_rms"]
+        step = state["step"]
+
+        exp_avg_sq = state["exp_avg_sq"]
+        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=(1 - beta2))
+
+        this_step = state["step"] - (state["zero_step"]
+                                     if "zero_step" in state else 0)
+        bias_correction2 = 1 - beta2**(this_step + 1)
+        if bias_correction2 < 0.99:
+            # note: not in-place.
+            exp_avg_sq = exp_avg_sq * (1.0 / bias_correction2)
+
+        denom = exp_avg_sq.sqrt()
+        denom += eps
+        grad = grad / denom
+
+        alpha = -lr * (1 - beta1) * state["param_rms"].clamp(min=param_min_rms)
+
+        delta = state["delta"]
+        delta.add_(grad * alpha)
+        p.add_(delta)
+
+    def _step_scalar(self, group: dict, p: Tensor, state: dict):
+        """
+        A simplified form of the core update for scalar tensors, where we cannot get a good
+        estimate of the parameter rms.
+        """
+        beta1, beta2 = group["betas"]
+        scalar_max = group["scalar_max"]
+        eps = group["eps"]
+        lr = group["lr"] * group["scalar_lr_scale"]
+        grad = p.grad
+
+        exp_avg_sq = state["exp_avg_sq"]  # shape: (batch_size,)
+        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+        # bias_correction2 is like in Adam.  Don't bother with bias_correction1;
+        # slower update at the start will help stability anyway.
+        bias_correction2 = 1 - beta2**(state["step"] + 1)
+        denom = (exp_avg_sq / bias_correction2).sqrt() + eps
+
+        delta = state["delta"]
+        delta.add_(grad / denom, alpha=-lr * (1 - beta1))
+        p.clamp_(min=-scalar_max, max=scalar_max)
+        p.add_(delta)

AR/modules/patched_mha_with_cache.py

@@ -0,0 +1,465 @@
+from torch.nn.functional import *
+from torch.nn.functional import (
+    _mha_shape_check,
+    _canonical_mask,
+    _none_or_dtype,
+    _in_projection_packed,
+)
+from torch.nn import functional as F
+import torch
+# Tensor = torch.Tensor
+# from typing import Callable, List, Optional, Tuple, Union
+
+
+def multi_head_attention_forward_patched(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    embed_dim_to_check: int,
+    num_heads: int,
+    in_proj_weight: Optional[Tensor],
+    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]]:
+    r"""
+    Args:
+        query, key, value: map a query and a set of key-value pairs to an output.
+            See "Attention Is All You Need" for more details.
+        embed_dim_to_check: total dimension of the model.
+        num_heads: parallel attention heads.
+        in_proj_weight, in_proj_bias: input projection weight and bias.
+        bias_k, bias_v: bias of the key and value sequences to be added at dim=0.
+        add_zero_attn: add a new batch of zeros to the key and
+                       value sequences at dim=1.
+        dropout_p: probability of an element to be zeroed.
+        out_proj_weight, out_proj_bias: the output projection weight and bias.
+        training: apply dropout if is ``True``.
+        key_padding_mask: if provided, specified padding elements in the key will
+            be ignored by the attention. This is an binary mask. When the value is True,
+            the corresponding value on the attention layer will be filled with -inf.
+        need_weights: output attn_output_weights.
+            Default: `True`
+            Note: `needs_weight` defaults to `True`, but should be set to `False`
+            For best performance when attention weights are not nedeeded.
+            *Setting needs_weights to `True`
+            leads to a significant performance degradation.*
+        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
+            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
+        is_causal: If specified, applies a causal mask as attention mask, and ignores
+            attn_mask for computing scaled dot product attention.
+            Default: ``False``.
+            .. warning::
+                is_causal is provides a hint that the attn_mask is the
+                causal mask.Providing incorrect hints can result in
+                incorrect execution, including forward and backward
+                compatibility.
+        use_separate_proj_weight: the function accept the proj. weights for query, key,
+            and value in different forms. If false, in_proj_weight will be used, which is
+            a combination of q_proj_weight, k_proj_weight, v_proj_weight.
+        q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.
+        static_k, static_v: static key and value used for attention operators.
+        average_attn_weights: If true, indicates that the returned ``attn_weights`` should be averaged across heads.
+            Otherwise, ``attn_weights`` are provided separately per head. Note that this flag only has an effect
+            when ``need_weights=True.``. Default: True
+
+
+    Shape:
+        Inputs:
+        - query: :math:`(L, E)` or :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key: :math:`(S, E)` or :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - value: :math:`(S, E)` or :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key_padding_mask: :math:`(S)` or :math:`(N, S)` where N is the batch size, S is the source sequence length.
+          If a FloatTensor is provided, it will be directly added to the value.
+          If a BoolTensor is provided, the positions with the
+          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
+        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
+          3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
+          S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked
+          positions. If a BoolTensor is provided, positions with ``True``
+          are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
+          is provided, it will be added to the attention weight.
+        - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
+          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
+        - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
+          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
+
+        Outputs:
+        - attn_output: :math:`(L, E)` or :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
+          E is the embedding dimension.
+        - attn_output_weights: Only returned when ``need_weights=True``. If ``average_attn_weights=True``, returns
+          attention weights averaged across heads of shape :math:`(L, S)` when input is unbatched or
+          :math:`(N, L, S)`, where :math:`N` is the batch size, :math:`L` is the target sequence length, and
+          :math:`S` is the source sequence length. If ``average_attn_weights=False``, returns attention weights per
+          head of shape :math:`(num_heads, L, S)` when input is unbatched or :math:`(N, num_heads, L, S)`.
+    """
+    tens_ops = (
+        query,
+        key,
+        value,
+        in_proj_weight,
+        in_proj_bias,
+        bias_k,
+        bias_v,
+        out_proj_weight,
+        out_proj_bias,
+    )
+    if has_torch_function(tens_ops):
+        return handle_torch_function(
+            multi_head_attention_forward,
+            tens_ops,
+            query,
+            key,
+            value,
+            embed_dim_to_check,
+            num_heads,
+            in_proj_weight,
+            in_proj_bias,
+            bias_k,
+            bias_v,
+            add_zero_attn,
+            dropout_p,
+            out_proj_weight,
+            out_proj_bias,
+            training=training,
+            key_padding_mask=key_padding_mask,
+            need_weights=need_weights,
+            attn_mask=attn_mask,
+            is_causal=is_causal,
+            use_separate_proj_weight=use_separate_proj_weight,
+            q_proj_weight=q_proj_weight,
+            k_proj_weight=k_proj_weight,
+            v_proj_weight=v_proj_weight,
+            static_k=static_k,
+            static_v=static_v,
+            average_attn_weights=average_attn_weights,
+            cache=cache,
+        )
+
+    is_batched = _mha_shape_check(
+        query, key, value, key_padding_mask, attn_mask, num_heads
+    )
+
+    # For unbatched input, we unsqueeze at the expected batch-dim to pretend that the input
+    # is batched, run the computation and before returning squeeze the
+    # batch dimension so that the output doesn't carry this temporary batch dimension.
+    if not is_batched:
+        # unsqueeze if the input is unbatched
+        query = query.unsqueeze(1)
+        key = key.unsqueeze(1)
+        value = value.unsqueeze(1)
+        if key_padding_mask is not None:
+            key_padding_mask = key_padding_mask.unsqueeze(0)
+
+    # set up shape vars
+    tgt_len, bsz, embed_dim = query.shape
+    src_len, _, _ = key.shape
+
+    key_padding_mask = _canonical_mask(
+        mask=key_padding_mask,
+        mask_name="key_padding_mask",
+        other_type=_none_or_dtype(attn_mask),
+        other_name="attn_mask",
+        target_type=query.dtype,
+    )
+
+    if is_causal and attn_mask is None:
+        raise RuntimeError(
+            "Need attn_mask if specifying the is_causal hint. "
+            "You may use the Transformer module method "
+            "`generate_square_subsequent_mask` to create this mask."
+        )
+
+    if is_causal and key_padding_mask is None and not need_weights:
+        # when we have a kpm or need weights, we need attn_mask
+        # Otherwise, we use the is_causal hint go as is_causal
+        # indicator to SDPA.
+        attn_mask = None
+    else:
+        attn_mask = _canonical_mask(
+            mask=attn_mask,
+            mask_name="attn_mask",
+            other_type=None,
+            other_name="",
+            target_type=query.dtype,
+            check_other=False,
+        )
+
+        if key_padding_mask is not None:
+            # We have the attn_mask, and use that to merge kpm into it.
+            # Turn off use of is_causal hint, as the merged mask is no
+            # longer causal.
+            is_causal = False
+
+    assert (
+        embed_dim == embed_dim_to_check
+    ), f"was expecting embedding dimension of {embed_dim_to_check}, but got {embed_dim}"
+    if isinstance(embed_dim, torch.Tensor):
+        # embed_dim can be a tensor when JIT tracing
+        head_dim = embed_dim.div(num_heads, rounding_mode="trunc")
+    else:
+        head_dim = embed_dim // num_heads
+    assert (
+        head_dim * num_heads == embed_dim
+    ), f"embed_dim {embed_dim} not divisible by num_heads {num_heads}"
+    if use_separate_proj_weight:
+        # allow MHA to have different embedding dimensions when separate projection weights are used
+        assert (
+            key.shape[:2] == value.shape[:2]
+        ), f"key's sequence and batch dims {key.shape[:2]} do not match value's {value.shape[:2]}"
+    else:
+        assert (
+            key.shape == value.shape
+        ), f"key shape {key.shape} does not match value shape {value.shape}"
+
+    #
+    # compute in-projection
+    #
+    if not use_separate_proj_weight:
+        assert (
+            in_proj_weight is not None
+        ), "use_separate_proj_weight is False but in_proj_weight is None"
+        q, k, v = _in_projection_packed(query, key, value, in_proj_weight, in_proj_bias)
+    else:
+        assert (
+            q_proj_weight is not None
+        ), "use_separate_proj_weight is True but q_proj_weight is None"
+        assert (
+            k_proj_weight is not None
+        ), "use_separate_proj_weight is True but k_proj_weight is None"
+        assert (
+            v_proj_weight is not None
+        ), "use_separate_proj_weight is True but v_proj_weight is None"
+        if in_proj_bias is None:
+            b_q = b_k = b_v = None
+        else:
+            b_q, b_k, b_v = in_proj_bias.chunk(3)
+        q, k, v = _in_projection(
+            query,
+            key,
+            value,
+            q_proj_weight,
+            k_proj_weight,
+            v_proj_weight,
+            b_q,
+            b_k,
+            b_v,
+        )
+    if cache != None:
+        if cache["first_infer"] == 1:
+            cache["k"][cache["stage"]] = k
+            # print(0,cache["k"].shape)
+            cache["v"][cache["stage"]] = v
+        else:  ###12个layer每个都要留自己的cache_kv
+            # print(1,cache["k"].shape)
+            cache["k"][cache["stage"]] = torch.cat(
+                [cache["k"][cache["stage"]], k], 0
+            )  ##本来时序是1，但是proj的时候可能transpose了所以时序到0维了
+            cache["v"][cache["stage"]] = torch.cat([cache["v"][cache["stage"]], v], 0)
+            # print(2, cache["k"].shape)
+            src_len = cache["k"][cache["stage"]].shape[0]
+            k = cache["k"][cache["stage"]]
+            v = cache["v"][cache["stage"]]
+            # if attn_mask is not None:
+            #     attn_mask=attn_mask[-1:,]
+            # print(attn_mask.shape,attn_mask)
+        cache["stage"] = (cache["stage"] + 1) % cache["all_stage"]
+    # print(2333,cache)
+    # prep attention mask
+
+    attn_mask = _canonical_mask(
+        mask=attn_mask,
+        mask_name="attn_mask",
+        other_type=None,
+        other_name="",
+        target_type=q.dtype,
+        check_other=False,
+    )
+
+    if attn_mask is not None:
+        # ensure attn_mask's dim is 3
+        if attn_mask.dim() == 2:
+            correct_2d_size = (tgt_len, src_len)
+            if attn_mask.shape != correct_2d_size:
+                raise RuntimeError(
+                    f"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}."
+                )
+            attn_mask = attn_mask.unsqueeze(0)
+        elif attn_mask.dim() == 3:
+            correct_3d_size = (bsz * num_heads, tgt_len, src_len)
+            if attn_mask.shape != correct_3d_size:
+                raise RuntimeError(
+                    f"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}."
+                )
+        else:
+            raise RuntimeError(
+                f"attn_mask's dimension {attn_mask.dim()} is not supported"
+            )
+
+    # add bias along batch dimension (currently second)
+    if bias_k is not None and bias_v is not None:
+        assert static_k is None, "bias cannot be added to static key."
+        assert static_v is None, "bias cannot be added to static value."
+        k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
+        v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
+        if attn_mask is not None:
+            attn_mask = pad(attn_mask, (0, 1))
+        if key_padding_mask is not None:
+            key_padding_mask = pad(key_padding_mask, (0, 1))
+    else:
+        assert bias_k is None
+        assert bias_v is None
+
+    #
+    # reshape q, k, v for multihead attention and make em batch first
+    #
+    q = q.view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
+    if static_k is None:
+        k = k.view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
+    else:
+        # TODO finish disentangling control flow so we don't do in-projections when statics are passed
+        assert (
+            static_k.size(0) == bsz * num_heads
+        ), f"expecting static_k.size(0) of {bsz * num_heads}, but got {static_k.size(0)}"
+        assert (
+            static_k.size(2) == head_dim
+        ), f"expecting static_k.size(2) of {head_dim}, but got {static_k.size(2)}"
+        k = static_k
+    if static_v is None:
+        v = v.view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
+    else:
+        # TODO finish disentangling control flow so we don't do in-projections when statics are passed
+        assert (
+            static_v.size(0) == bsz * num_heads
+        ), f"expecting static_v.size(0) of {bsz * num_heads}, but got {static_v.size(0)}"
+        assert (
+            static_v.size(2) == head_dim
+        ), f"expecting static_v.size(2) of {head_dim}, but got {static_v.size(2)}"
+        v = static_v
+
+    # add zero attention along batch dimension (now first)
+    if add_zero_attn:
+        zero_attn_shape = (bsz * num_heads, 1, head_dim)
+        k = torch.cat(
+            [k, torch.zeros(zero_attn_shape, dtype=k.dtype, device=k.device)], dim=1
+        )
+        v = torch.cat(
+            [v, torch.zeros(zero_attn_shape, dtype=v.dtype, device=v.device)], dim=1
+        )
+        if attn_mask is not None:
+            attn_mask = pad(attn_mask, (0, 1))
+        if key_padding_mask is not None:
+            key_padding_mask = pad(key_padding_mask, (0, 1))
+
+    # update source sequence length after adjustments
+    src_len = k.size(1)
+
+    # merge key padding and attention masks
+    if key_padding_mask is not None:
+        assert key_padding_mask.shape == (
+            bsz,
+            src_len,
+        ), f"expecting key_padding_mask shape of {(bsz, src_len)}, but got {key_padding_mask.shape}"
+        key_padding_mask = (
+            key_padding_mask.view(bsz, 1, 1, src_len)
+            .expand(-1, num_heads, -1, -1)
+            .reshape(bsz * num_heads, 1, src_len)
+        )
+        if attn_mask is None:
+            attn_mask = key_padding_mask
+        else:
+            attn_mask = attn_mask + key_padding_mask
+
+    # adjust dropout probability
+    if not training:
+        dropout_p = 0.0
+
+    #
+    # (deep breath) calculate attention and out projection
+    #
+
+    if need_weights:
+        B, Nt, E = q.shape
+        q_scaled = q / math.sqrt(E)
+
+        assert not (
+            is_causal and attn_mask is None
+        ), "FIXME: is_causal not implemented for need_weights"
+
+        if attn_mask is not None:
+            attn_output_weights = torch.baddbmm(
+                attn_mask, q_scaled, k.transpose(-2, -1)
+            )
+        else:
+            attn_output_weights = torch.bmm(q_scaled, k.transpose(-2, -1))
+        attn_output_weights = softmax(attn_output_weights, dim=-1)
+        if dropout_p > 0.0:
+            attn_output_weights = dropout(attn_output_weights, p=dropout_p)
+
+        attn_output = torch.bmm(attn_output_weights, v)
+
+        attn_output = (
+            attn_output.transpose(0, 1).contiguous().view(tgt_len * bsz, embed_dim)
+        )
+        attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
+        attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
+
+        # optionally average attention weights over heads
+        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
+        if average_attn_weights:
+            attn_output_weights = attn_output_weights.mean(dim=1)
+
+        if not is_batched:
+            # squeeze the output if input was unbatched
+            attn_output = attn_output.squeeze(1)
+            attn_output_weights = attn_output_weights.squeeze(0)
+        return attn_output, attn_output_weights
+    else:
+        # attn_mask can be either (L,S) or (N*num_heads, L, S)
+        # if attn_mask's shape is (1, L, S) we need to unsqueeze to (1, 1, L, S)
+        # in order to match the input for SDPA of (N, num_heads, L, S)
+        if attn_mask is not None:
+            if attn_mask.size(0) == 1 and attn_mask.dim() == 3:
+                attn_mask = attn_mask.unsqueeze(0)
+            else:
+                attn_mask = attn_mask.view(bsz, num_heads, -1, src_len)
+
+        q = q.view(bsz, num_heads, tgt_len, head_dim)
+        k = k.view(bsz, num_heads, src_len, head_dim)
+        v = v.view(bsz, num_heads, src_len, head_dim)
+
+        # with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True):
+        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(bsz * tgt_len, embed_dim)
+        )
+
+        attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
+        attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
+        if not is_batched:
+            # squeeze the output if input was unbatched
+            attn_output = attn_output.squeeze(1)
+        return attn_output, None

AR/modules/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

AR/modules/scaling.py

@@ -0,0 +1,335 @@
+# Copyright    2022  Xiaomi Corp.        (authors: Daniel Povey)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# 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 logging
+import math
+import random
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+class DoubleSwishFunction(torch.autograd.Function):
+    """
+      double_swish(x) = x * torch.sigmoid(x-1)
+    This is a definition, originally motivated by its close numerical
+    similarity to swish(swish(x)), where swish(x) =  x * sigmoid(x).
+
+    Memory-efficient derivative computation:
+     double_swish(x) = x * s, where s(x) = torch.sigmoid(x-1)
+     double_swish'(x) = d/dx double_swish(x) =  x * s'(x) + x' * s(x) = x * s'(x) + s(x).
+     Now, s'(x) = s(x) * (1-s(x)).
+     double_swish'(x) =  x * s'(x) + s(x).
+                      =  x * s(x) * (1-s(x)) + s(x).
+                     = double_swish(x) * (1-s(x)) + s(x)
+     ... so we just need to remember s(x) but not x itself.
+    """
+
+    @staticmethod
+    def forward(ctx, x: Tensor) -> Tensor:
+        requires_grad = x.requires_grad
+        x_dtype = x.dtype
+        if x.dtype == torch.float16:
+            x = x.to(torch.float32)
+
+        s = torch.sigmoid(x - 1.0)
+        y = x * s
+
+        if requires_grad:
+            deriv = y * (1 - s) + s
+            # notes on derivative of x * sigmoid(x - 1):
+            # https://www.wolframalpha.com/input?i=d%2Fdx+%28x+*+sigmoid%28x-1%29%29
+            # min \simeq -0.043638.  Take floor as -0.043637 so it's a lower bund
+            # max \simeq 1.1990.   Take ceil to be 1.2 so it's an upper bound.
+            # the combination of "+ torch.rand_like(deriv)" and casting to torch.uint8 (which
+            # floors), should be expectation-preserving.
+            floor = -0.043637
+            ceil = 1.2
+            d_scaled = (deriv - floor) * (255.0 / (ceil - floor)) + torch.rand_like(
+                deriv
+            )
+            if __name__ == "__main__":
+                # for self-testing only.
+                assert d_scaled.min() >= 0.0
+                assert d_scaled.max() < 256.0
+            d_int = d_scaled.to(torch.uint8)
+            ctx.save_for_backward(d_int)
+        if x.dtype == torch.float16 or torch.is_autocast_enabled():
+            y = y.to(torch.float16)
+        return y
+
+    @staticmethod
+    def backward(ctx, y_grad: Tensor) -> Tensor:
+        (d,) = ctx.saved_tensors
+        # the same constants as used in forward pass.
+        floor = -0.043637
+        ceil = 1.2
+        d = d * ((ceil - floor) / 255.0) + floor
+        return y_grad * d
+
+
+class DoubleSwish(torch.nn.Module):
+    def forward(self, x: Tensor) -> Tensor:
+        """Return double-swish activation function which is an approximation to Swish(Swish(x)),
+        that we approximate closely with x * sigmoid(x-1).
+        """
+        if torch.jit.is_scripting() or torch.jit.is_tracing():
+            return x * torch.sigmoid(x - 1.0)
+        return DoubleSwishFunction.apply(x)
+
+
+class ActivationBalancerFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        scale_factor: Tensor,
+        sign_factor: Optional[Tensor],
+        channel_dim: int,
+    ) -> Tensor:
+        if channel_dim < 0:
+            channel_dim += x.ndim
+        ctx.channel_dim = channel_dim
+        xgt0 = x > 0
+        if sign_factor is None:
+            ctx.save_for_backward(xgt0, scale_factor)
+        else:
+            ctx.save_for_backward(xgt0, scale_factor, sign_factor)
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad: Tensor) -> Tuple[Tensor, None, None, None]:
+        if len(ctx.saved_tensors) == 3:
+            xgt0, scale_factor, sign_factor = ctx.saved_tensors
+            for _ in range(ctx.channel_dim, x_grad.ndim - 1):
+                scale_factor = scale_factor.unsqueeze(-1)
+                sign_factor = sign_factor.unsqueeze(-1)
+            factor = sign_factor + scale_factor * (xgt0.to(x_grad.dtype) - 0.5)
+        else:
+            xgt0, scale_factor = ctx.saved_tensors
+            for _ in range(ctx.channel_dim, x_grad.ndim - 1):
+                scale_factor = scale_factor.unsqueeze(-1)
+            factor = scale_factor * (xgt0.to(x_grad.dtype) - 0.5)
+        neg_delta_grad = x_grad.abs() * factor
+        return (
+            x_grad - neg_delta_grad,
+            None,
+            None,
+            None,
+        )
+
+
+def _compute_scale_factor(
+    x: Tensor,
+    channel_dim: int,
+    min_abs: float,
+    max_abs: float,
+    gain_factor: float,
+    max_factor: float,
+) -> Tensor:
+    if channel_dim < 0:
+        channel_dim += x.ndim
+    sum_dims = [d for d in range(x.ndim) if d != channel_dim]
+    x_abs_mean = torch.mean(x.abs(), dim=sum_dims).to(torch.float32)
+
+    if min_abs == 0.0:
+        below_threshold = 0.0
+    else:
+        # below_threshold is 0 if x_abs_mean > min_abs, can be at most max_factor if
+        # x_abs)_mean , min_abs.
+        below_threshold = ((min_abs - x_abs_mean) * (gain_factor / min_abs)).clamp(
+            min=0, max=max_factor
+        )
+
+    above_threshold = ((x_abs_mean - max_abs) * (gain_factor / max_abs)).clamp(
+        min=0, max=max_factor
+    )
+
+    return below_threshold - above_threshold
+
+
+def _compute_sign_factor(
+    x: Tensor,
+    channel_dim: int,
+    min_positive: float,
+    max_positive: float,
+    gain_factor: float,
+    max_factor: float,
+) -> Tensor:
+    if channel_dim < 0:
+        channel_dim += x.ndim
+    sum_dims = [d for d in range(x.ndim) if d != channel_dim]
+    proportion_positive = torch.mean((x > 0).to(torch.float32), dim=sum_dims)
+    if min_positive == 0.0:
+        factor1 = 0.0
+    else:
+        # 0 if proportion_positive >= min_positive, else can be
+        # as large as max_factor.
+        factor1 = (
+            (min_positive - proportion_positive) * (gain_factor / min_positive)
+        ).clamp_(min=0, max=max_factor)
+
+    if max_positive == 1.0:
+        factor2 = 0.0
+    else:
+        # 0 if self.proportion_positive <= max_positive, else can be
+        # as large as -max_factor.
+        factor2 = (
+            (proportion_positive - max_positive) * (gain_factor / (1.0 - max_positive))
+        ).clamp_(min=0, max=max_factor)
+    sign_factor = factor1 - factor2
+    # require min_positive != 0 or max_positive != 1:
+    assert not isinstance(sign_factor, float)
+    return sign_factor
+
+
+class ActivationBalancer(torch.nn.Module):
+    """
+    Modifies the backpropped derivatives of a function to try to encourage, for
+    each channel, that it is positive at least a proportion `threshold` of the
+    time.  It does this by multiplying negative derivative values by up to
+    (1+max_factor), and positive derivative values by up to (1-max_factor),
+    interpolated from 1 at the threshold to those extremal values when none
+    of the inputs are positive.
+
+    Args:
+           num_channels: the number of channels
+           channel_dim: the dimension/axis corresponding to the channel, e.g.
+               -1, 0, 1, 2; will be interpreted as an offset from x.ndim if negative.
+           min_positive: the minimum, per channel, of the proportion of the time
+               that (x > 0), below which we start to modify the derivatives.
+           max_positive: the maximum, per channel, of the proportion of the time
+               that (x > 0), above which we start to modify the derivatives.
+           max_factor: the maximum factor by which we modify the derivatives for
+              either the sign constraint or the magnitude constraint;
+              e.g. with max_factor=0.02, the the derivatives would be multiplied by
+              values in the range [0.98..1.02].
+           sign_gain_factor: determines the 'gain' with which we increase the
+              change in gradient once the constraints on min_positive and max_positive
+              are violated.
+           scale_gain_factor: determines the 'gain' with which we increase the
+              change in gradient once the constraints on min_abs and max_abs
+              are violated.
+           min_abs:  the minimum average-absolute-value difference from the mean
+               value per channel, which we allow, before we start to modify
+               the derivatives to prevent this.
+           max_abs:  the maximum average-absolute-value difference from the mean
+               value per channel, which we allow, before we start to modify
+               the derivatives to prevent this.
+          min_prob: determines the minimum probability with which we modify the
+             gradients for the {min,max}_positive and {min,max}_abs constraints,
+             on each forward().  This is done randomly to prevent all layers
+             from doing it at the same time.  Early in training we may use
+             higher probabilities than this; it will decay to this value.
+    """
+
+    def __init__(
+        self,
+        num_channels: int,
+        channel_dim: int,
+        min_positive: float = 0.05,
+        max_positive: float = 0.95,
+        max_factor: float = 0.04,
+        sign_gain_factor: float = 0.01,
+        scale_gain_factor: float = 0.02,
+        min_abs: float = 0.2,
+        max_abs: float = 100.0,
+        min_prob: float = 0.1,
+    ):
+        super(ActivationBalancer, self).__init__()
+        self.num_channels = num_channels
+        self.channel_dim = channel_dim
+        self.min_positive = min_positive
+        self.max_positive = max_positive
+        self.max_factor = max_factor
+        self.min_abs = min_abs
+        self.max_abs = max_abs
+        self.min_prob = min_prob
+        self.sign_gain_factor = sign_gain_factor
+        self.scale_gain_factor = scale_gain_factor
+
+        # count measures how many times the forward() function has been called.
+        # We occasionally sync this to a tensor called `count`, that exists to
+        # make sure it is synced to disk when we load and save the model.
+        self.cpu_count = 0
+        self.register_buffer("count", torch.tensor(0, dtype=torch.int64))
+
+    def forward(self, x: Tensor) -> Tensor:
+        if torch.jit.is_scripting() or not x.requires_grad or torch.jit.is_tracing():
+            return _no_op(x)
+
+        count = self.cpu_count
+        self.cpu_count += 1
+
+        if random.random() < 0.01:
+            # Occasionally sync self.cpu_count with self.count.
+            # count affects the decay of 'prob'.  don't do this on every iter,
+            # because syncing with the GPU is slow.
+            self.cpu_count = max(self.cpu_count, self.count.item())
+            self.count.fill_(self.cpu_count)
+
+        # the prob of doing some work exponentially decreases from 0.5 till it hits
+        # a floor at min_prob (==0.1, by default)
+        prob = max(self.min_prob, 0.5 ** (1 + (count / 4000.0)))
+
+        if random.random() < prob:
+            sign_gain_factor = 0.5
+            if self.min_positive != 0.0 or self.max_positive != 1.0:
+                sign_factor = _compute_sign_factor(
+                    x,
+                    self.channel_dim,
+                    self.min_positive,
+                    self.max_positive,
+                    gain_factor=self.sign_gain_factor / prob,
+                    max_factor=self.max_factor,
+                )
+            else:
+                sign_factor = None
+
+            scale_factor = _compute_scale_factor(
+                x.detach(),
+                self.channel_dim,
+                min_abs=self.min_abs,
+                max_abs=self.max_abs,
+                gain_factor=self.scale_gain_factor / prob,
+                max_factor=self.max_factor,
+            )
+            return ActivationBalancerFunction.apply(
+                x,
+                scale_factor,
+                sign_factor,
+                self.channel_dim,
+            )
+        else:
+            return _no_op(x)
+
+
+def BalancedDoubleSwish(
+    d_model, channel_dim=-1, max_abs=10.0, min_prob=0.25
+) -> nn.Sequential:
+    """
+    ActivationBalancer -> DoubleSwish
+    """
+    balancer = ActivationBalancer(
+        d_model, channel_dim=channel_dim, max_abs=max_abs, min_prob=min_prob
+    )
+    return nn.Sequential(
+        balancer,
+        DoubleSwish(),
+    )

AR/modules/transformer.py

@@ -0,0 +1,379 @@
+# 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 torch import Tensor
+from torch import nn
+from torch.nn import functional as F
+
+from moyoyo_tts.AR.modules.activation import MultiheadAttention
+from moyoyo_tts.AR.modules.scaling import BalancedDoubleSwish
+
+_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:
+        r"""Pass the input through the encoder layers in turn.
+
+        Args:
+            src: the sequence to the encoder (required).
+            mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+            return_layer_states: return layers' state (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        if return_layer_states:
+            layer_states = []  # layers' output
+            output = src
+            for mod in self.layers:
+                output = mod(
+                    output,
+                    src_mask=mask,
+                    src_key_padding_mask=src_key_padding_mask,
+                    cache=cache,
+                )
+                layer_states.append(output[0])
+
+            if self.norm is not None:
+                output = self.norm(output)
+
+            return layer_states, output
+
+        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__()
+        # print(233333333333,d_model,nhead)
+        # import os
+        # os._exit(2333333)
+        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,
+        )
+
+        # Implementation of Feedforward model
+        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)
+
+        # Legacy string support for activation function.
+        if isinstance(activation, str):
+            activation = _get_activation_fn(activation)
+        elif isinstance(activation, partial):
+            activation = activation(d_model)
+        elif activation == BalancedDoubleSwish:
+            activation = BalancedDoubleSwish(d_model)
+
+        # # We can't test self.activation in forward() in TorchScript,
+        # # so stash some information about it instead.
+        # if activation is F.relu or isinstance(activation, torch.nn.ReLU):
+        #     self.activation_relu_or_gelu = 1
+        # elif activation is F.gelu or isinstance(activation, torch.nn.GELU):
+        #     self.activation_relu_or_gelu = 2
+        # else:
+        #     self.activation_relu_or_gelu = 0
+        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:
+        r"""Pass the input through the encoder layer.
+
+        Args:
+            src: the sequence to the encoder layer (required).
+            src_mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        x, stage_embedding = src, None
+        is_src_tuple = False
+        if isinstance(src, tuple):
+            x, stage_embedding = src
+            is_src_tuple = True
+
+        if src_key_padding_mask is not None:
+            _skpm_dtype = src_key_padding_mask.dtype
+            if _skpm_dtype != torch.bool and not torch.is_floating_point(
+                src_key_padding_mask
+            ):
+                raise AssertionError(
+                    "only bool and floating types of key_padding_mask are supported"
+                )
+
+        if self.norm_first:
+            x = x + self._sa_block(
+                self.norm1(x, stage_embedding),
+                src_mask,
+                src_key_padding_mask,
+                cache=cache,
+            )
+            x = x + self._ff_block(self.norm2(x, stage_embedding))
+        else:
+            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)
+
+        if is_src_tuple:
+            return (x, stage_embedding)
+        return x
+
+    # self-attention block
+    def _sa_block(
+        self,
+        x: Tensor,
+        attn_mask: Optional[Tensor],
+        key_padding_mask: Optional[Tensor],
+        cache=None,
+    ) -> Tensor:
+        # print(x.shape,attn_mask.shape,key_padding_mask)
+        # torch.Size([1, 188, 512]) torch.Size([188, 188]) None
+        # import os
+        # os._exit(23333)
+        x = self.self_attn(
+            x,
+            x,
+            x,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=False,
+            cache=cache,
+        )[0]
+        return self.dropout1(x)
+
+    # feed forward block
+    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)])

AR/modules/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 moyoyo_tts.AR.modules.activation_onnx import MultiheadAttention
+from moyoyo_tts.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)])

AR/text_processing/phonemizer.py

@@ -0,0 +1,79 @@
+# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/text_processing/phonemizer.py
+# reference: https://github.com/lifeiteng/vall-e
+import itertools
+import re
+from typing import Dict
+from typing import List
+
+import regex
+from gruut import sentences
+from gruut.const import Sentence
+from gruut.const import Word
+from moyoyo_tts.AR.text_processing.symbols import SYMBOL_TO_ID
+
+
+class GruutPhonemizer:
+    def __init__(self, language: str):
+        self._phonemizer = sentences
+        self.lang = language
+        self.symbol_to_id = SYMBOL_TO_ID
+        self._special_cases_dict: Dict[str] = {
+            r"\.\.\.": "... ",
+            ";": "; ",
+            ":": ": ",
+            ",": ", ",
+            r"\.": ". ",
+            "!": "! ",
+            r"\?": "? ",
+            "—": "—",
+            "…": "… ",
+            "«": "«",
+            "»": "»",
+        }
+        self._punctuation_regexp: str = (
+            rf"([{''.join(self._special_cases_dict.keys())}])"
+        )
+
+    def _normalize_punctuation(self, text: str) -> str:
+        text = regex.sub(rf"\pZ+{self._punctuation_regexp}", r"\1", text)
+        text = regex.sub(rf"{self._punctuation_regexp}(\pL)", r"\1 \2", text)
+        text = regex.sub(r"\pZ+", r" ", text)
+        return text.strip()
+
+    def _convert_punctuation(self, word: Word) -> str:
+        if not word.phonemes:
+            return ""
+        if word.phonemes[0] in ["‖", "|"]:
+            return word.text.strip()
+
+        phonemes = "".join(word.phonemes)
+        # remove modifier characters ˈˌː with regex
+        phonemes = re.sub(r"[ˈˌː͡]", "", phonemes)
+        return phonemes.strip()
+
+    def phonemize(self, text: str, espeak: bool = False) -> str:
+        text_to_phonemize: str = self._normalize_punctuation(text)
+        sents: List[Sentence] = [
+            sent
+            for sent in self._phonemizer(text_to_phonemize, lang="en-us", espeak=espeak)
+        ]
+        words: List[str] = [
+            self._convert_punctuation(word) for word in itertools.chain(*sents)
+        ]
+        return " ".join(words)
+
+    def transform(self, phonemes):
+        # convert phonemes to ids
+        # dictionary is in symbols.py
+        return [self.symbol_to_id[p] for p in phonemes if p in self.symbol_to_id.keys()]
+
+
+if __name__ == "__main__":
+    phonemizer = GruutPhonemizer("en-us")
+    # text -> IPA
+    phonemes = phonemizer.phonemize("Hello, wor-ld ?")
+    print("phonemes:", phonemes)
+    print("len(phonemes):", len(phonemes))
+    phoneme_ids = phonemizer.transform(phonemes)
+    print("phoneme_ids:", phoneme_ids)
+    print("len(phoneme_ids):", len(phoneme_ids))

AR/text_processing/symbols.py

@@ -0,0 +1,10 @@
+# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/text_processing/symbols.py
+# reference: https://github.com/lifeiteng/vall-e
+PAD = "_"
+PUNCTUATION = ';:,.!?¡¿—…"«»“” '
+LETTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+IPA_LETTERS = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
+SYMBOLS = [PAD] + list(PUNCTUATION) + list(LETTERS) + list(IPA_LETTERS)
+SPACE_ID = SYMBOLS.index(" ")
+SYMBOL_TO_ID = {s: i for i, s in enumerate(SYMBOLS)}
+ID_TO_SYMBOL = {i: s for i, s in enumerate(SYMBOLS)}

AR/utils/__init__.py

@@ -0,0 +1,37 @@
+import re
+
+
+def str2bool(str):
+    return True if str.lower() == 'true' else False
+
+
+def get_newest_ckpt(string_list):
+    # 定义一个正则表达式模式，用于匹配字符串中的数字
+    pattern = r'epoch=(\d+)-step=(\d+)\.ckpt'
+
+    # 使用正则表达式提取每个字符串中的数字信息，并创建一个包含元组的列表
+    extracted_info = []
+    for string in string_list:
+        match = re.match(pattern, string)
+        if match:
+            epoch = int(match.group(1))
+            step = int(match.group(2))
+            extracted_info.append((epoch, step, string))
+    # 按照 epoch 后面的数字和 step 后面的数字进行排序
+    sorted_info = sorted(
+        extracted_info, key=lambda x: (x[0], x[1]), reverse=True)
+    # 获取最新的 ckpt 文件名
+    newest_ckpt = sorted_info[0][2]
+    return newest_ckpt
+
+
+# 文本存在且不为空时 return True
+def check_txt_file(file_path):
+    try:
+        with open(file_path, 'r') as file:
+            text = file.readline().strip()
+        assert text.strip() != ''
+        return text
+    except Exception:
+        return False
+    return False

AR/utils/initialize.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+"""Initialize modules for espnet2 neural networks."""
+import torch
+from typeguard import check_argument_types
+
+
+def initialize(model: torch.nn.Module, init: str):
+    """Initialize weights of a neural network module.
+
+    Parameters are initialized using the given method or distribution.
+
+    Custom initialization routines can be implemented into submodules
+    as function `espnet_initialization_fn` within the custom module.
+
+    Args:
+        model: Target.
+        init: Method of initialization.
+    """
+    assert check_argument_types()
+    print("init with", init)
+
+    # weight init
+    for p in model.parameters():
+        if p.dim() > 1:
+            if init == "xavier_uniform":
+                torch.nn.init.xavier_uniform_(p.data)
+            elif init == "xavier_normal":
+                torch.nn.init.xavier_normal_(p.data)
+            elif init == "kaiming_uniform":
+                torch.nn.init.kaiming_uniform_(p.data, nonlinearity="relu")
+            elif init == "kaiming_normal":
+                torch.nn.init.kaiming_normal_(p.data, nonlinearity="relu")
+            else:
+                raise ValueError("Unknown initialization: " + init)
+    # bias init
+    for name, p in model.named_parameters():
+        if ".bias" in name and p.dim() == 1:
+            p.data.zero_()

AR/utils/io.py

@@ -0,0 +1,34 @@
+import sys
+
+import torch
+import yaml
+
+
+def load_yaml_config(path):
+    with open(path) as f:
+        config = yaml.full_load(f)
+    return config
+
+
+def save_config_to_yaml(config, path):
+    assert path.endswith(".yaml")
+    with open(path, "w") as f:
+        f.write(yaml.dump(config))
+        f.close()
+
+
+def write_args(args, path):
+    args_dict = dict(
+        (name, getattr(args, name)) for name in dir(args) if not name.startswith("_")
+    )
+    with open(path, "a") as args_file:
+        args_file.write("==> torch version: {}\n".format(torch.__version__))
+        args_file.write(
+            "==> cudnn version: {}\n".format(torch.backends.cudnn.version())
+        )
+        args_file.write("==> Cmd:\n")
+        args_file.write(str(sys.argv))
+        args_file.write("\n==> args:\n")
+        for k, v in sorted(args_dict.items()):
+            args_file.write("  %s: %s\n" % (str(k), str(v)))
+        args_file.close()

README.md

@@ -0,0 +1,348 @@
+<div align="center">
+
+
+<h1>GPT-SoVITS-WebUI</h1>
+A Powerful Few-shot Voice Conversion and Text-to-Speech WebUI.<br><br>
+
+[![madewithlove](https://img.shields.io/badge/made_with-%E2%9D%A4-red?style=for-the-badge&labelColor=orange)](https://github.com/RVC-Boss/GPT-SoVITS)
+
+<a href="https://trendshift.io/repositories/7033" target="_blank"><img src="https://trendshift.io/api/badge/repositories/7033" alt="RVC-Boss%2FGPT-SoVITS | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
+
+<!-- img src="https://counter.seku.su/cmoe?name=gptsovits&theme=r34" /><br> -->
+
+[![Open In Colab](https://img.shields.io/badge/Colab-F9AB00?style=for-the-badge&logo=googlecolab&color=525252)](https://colab.research.google.com/github/RVC-Boss/GPT-SoVITS/blob/main/colab_webui.ipynb)
+[![License](https://img.shields.io/badge/LICENSE-MIT-green.svg?style=for-the-badge)](https://github.com/RVC-Boss/GPT-SoVITS/blob/main/LICENSE)
+[![Huggingface](https://img.shields.io/badge/🤗%20-online%20demo-yellow.svg?style=for-the-badge)](https://huggingface.co/spaces/lj1995/GPT-SoVITS-v2)
+[![Discord](https://img.shields.io/discord/1198701940511617164?color=%23738ADB&label=Discord&style=for-the-badge)](https://discord.gg/dnrgs5GHfG)
+
+**English** | [**中文简体**](./docs/cn/README.md) | [**日本語**](./docs/ja/README.md) | [**한국어**](./docs/ko/README.md) | [**Türkçe**](./docs/tr/README.md)
+
+</div>
+
+---
+
+## Features:
+
+1. **Zero-shot TTS:** Input a 5-second vocal sample and experience instant text-to-speech conversion.
+
+2. **Few-shot TTS:** Fine-tune the model with just 1 minute of training data for improved voice similarity and realism.
+
+3. **Cross-lingual Support:** Inference in languages different from the training dataset, currently supporting English, Japanese, Korean, Cantonese and Chinese.
+
+4. **WebUI Tools:** Integrated tools include voice accompaniment separation, automatic training set segmentation, Chinese ASR, and text labeling, assisting beginners in creating training datasets and GPT/SoVITS models.
+
+**Check out our [demo video](https://www.bilibili.com/video/BV12g4y1m7Uw) here!**
+
+Unseen speakers few-shot fine-tuning demo:
+
+https://github.com/RVC-Boss/GPT-SoVITS/assets/129054828/05bee1fa-bdd8-4d85-9350-80c060ab47fb
+
+**User guide: [简体中文](https://www.yuque.com/baicaigongchang1145haoyuangong/ib3g1e) | [English](https://rentry.co/GPT-SoVITS-guide#/)**
+
+## Installation
+
+For users in China, you can [click here](https://www.codewithgpu.com/i/RVC-Boss/GPT-SoVITS/GPT-SoVITS-Official) to use AutoDL Cloud Docker to experience the full functionality online.
+
+### Tested Environments
+
+- Python 3.9, PyTorch 2.0.1, CUDA 11
+- Python 3.10.13, PyTorch 2.1.2, CUDA 12.3
+- Python 3.9, PyTorch 2.2.2, macOS 14.4.1 (Apple silicon)
+- Python 3.9, PyTorch 2.2.2, CPU devices
+
+_Note: numba==0.56.4 requires py<3.11_
+
+### Windows
+
+If you are a Windows user (tested with win>=10), you can [download the integrated package](https://huggingface.co/lj1995/GPT-SoVITS-windows-package/resolve/main/GPT-SoVITS-beta.7z?download=true) and double-click on _go-webui.bat_ to start GPT-SoVITS-WebUI.
+
+**Users in China can [download the package here](https://www.yuque.com/baicaigongchang1145haoyuangong/ib3g1e/dkxgpiy9zb96hob4#KTvnO).**
+
+### Linux
+
+```bash
+conda create -n GPTSoVits python=3.9
+conda activate GPTSoVits
+bash install.sh
+```
+
+### macOS
+
+**Note: The models trained with GPUs on Macs result in significantly lower quality compared to those trained on other devices, so we are temporarily using CPUs instead.**
+
+1. Install Xcode command-line tools by running `xcode-select --install`.
+2. Install FFmpeg by running `brew install ffmpeg`.
+3. Install the program by running the following commands:
+
+```bash
+conda create -n GPTSoVits python=3.9
+conda activate GPTSoVits
+pip install -r requirements.txt
+```
+
+### Install Manually
+
+#### Install FFmpeg
+
+##### Conda Users
+
+```bash
+conda install ffmpeg
+```
+
+##### Ubuntu/Debian Users
+
+```bash
+sudo apt install ffmpeg
+sudo apt install libsox-dev
+conda install -c conda-forge 'ffmpeg<7'
+```
+
+##### Windows Users
+
+Download and place [ffmpeg.exe](https://huggingface.co/lj1995/VoiceConversionWebUI/blob/main/ffmpeg.exe) and [ffprobe.exe](https://huggingface.co/lj1995/VoiceConversionWebUI/blob/main/ffprobe.exe) in the GPT-SoVITS root.
+
+Install [Visual Studio 2017](https://aka.ms/vs/17/release/vc_redist.x86.exe) (Korean TTS Only)
+
+##### MacOS Users
+```bash
+brew install ffmpeg
+```
+
+#### Install Dependences
+
+```bash
+pip install -r requirements.txt
+```
+
+### Using Docker
+
+#### docker-compose.yaml configuration
+
+0. Regarding image tags: Due to rapid updates in the codebase and the slow process of packaging and testing images, please check [Docker Hub](https://hub.docker.com/r/breakstring/gpt-sovits) for the currently packaged latest images and select as per your situation, or alternatively, build locally using a Dockerfile according to your own needs.
+1. Environment Variables：
+
+- is_half: Controls half-precision/double-precision. This is typically the cause if the content under the directories 4-cnhubert/5-wav32k is not generated correctly during the "SSL extracting" step. Adjust to True or False based on your actual situation.
+
+2. Volumes Configuration，The application's root directory inside the container is set to /workspace. The default docker-compose.yaml lists some practical examples for uploading/downloading content.
+3. shm_size： The default available memory for Docker Desktop on Windows is too small, which can cause abnormal operations. Adjust according to your own situation.
+4. Under the deploy section, GPU-related settings should be adjusted cautiously according to your system and actual circumstances.
+
+#### Running with docker compose
+
+```
+docker compose -f "docker-compose.yaml" up -d
+```
+
+#### Running with docker command
+
+As above, modify the corresponding parameters based on your actual situation, then run the following command:
+
+```
+docker run --rm -it --gpus=all --env=is_half=False --volume=G:\GPT-SoVITS-DockerTest\output:/workspace/output --volume=G:\GPT-SoVITS-DockerTest\logs:/workspace/logs --volume=G:\GPT-SoVITS-DockerTest\SoVITS_weights:/workspace/SoVITS_weights --workdir=/workspace -p 9880:9880 -p 9871:9871 -p 9872:9872 -p 9873:9873 -p 9874:9874 --shm-size="16G" -d breakstring/gpt-sovits:xxxxx
+```
+
+## Pretrained Models
+
+**Users in China can [download all these models here](https://www.yuque.com/baicaigongchang1145haoyuangong/ib3g1e/dkxgpiy9zb96hob4#nVNhX).**
+
+1. Download pretrained models from [GPT-SoVITS Models](https://huggingface.co/lj1995/GPT-SoVITS) and place them in `GPT_SoVITS/pretrained_models`.
+
+2. Download G2PW models from [G2PWModel_1.1.zip](https://paddlespeech.bj.bcebos.com/Parakeet/released_models/g2p/G2PWModel_1.1.zip), unzip and rename to `G2PWModel`, and then place them in `GPT_SoVITS/text`.(Chinese TTS Only)
+
+3. For UVR5 (Vocals/Accompaniment Separation & Reverberation Removal, additionally), download models from [UVR5 Weights](https://huggingface.co/lj1995/VoiceConversionWebUI/tree/main/uvr5_weights) and place them in `tools/uvr5/uvr5_weights`.
+
+4. For Chinese ASR (additionally), download models from [Damo ASR Model](https://modelscope.cn/models/damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/files), [Damo VAD Model](https://modelscope.cn/models/damo/speech_fsmn_vad_zh-cn-16k-common-pytorch/files), and [Damo Punc Model](https://modelscope.cn/models/damo/punc_ct-transformer_zh-cn-common-vocab272727-pytorch/files) and place them in `tools/asr/models`.
+
+5. For English or Japanese ASR (additionally), download models from [Faster Whisper Large V3](https://huggingface.co/Systran/faster-whisper-large-v3) and place them in `tools/asr/models`. Also, [other models](https://huggingface.co/Systran) may have the similar effect with smaller disk footprint. 
+
+## Dataset Format
+
+The TTS annotation .list file format:
+
+```
+vocal_path|speaker_name|language|text
+```
+
+Language dictionary:
+
+- 'zh': Chinese
+- 'ja': Japanese
+- 'en': English
+- 'ko': Korean
+- 'yue': Cantonese
+  
+Example:
+
+```
+D:\GPT-SoVITS\xxx/xxx.wav|xxx|en|I like playing Genshin.
+```
+
+## Finetune and inference
+
+ ### Open WebUI
+
+ #### Integrated Package Users
+
+ Double-click `go-webui.bat`or use `go-webui.ps1`
+ if you want to switch to V1,then double-click`go-webui-v1.bat` or use `go-webui-v1.ps1`
+
+ #### Others
+
+ ```bash
+ python webui.py <language(optional)>
+ ```
+
+ if you want to switch to V1,then
+
+ ```bash
+ python webui.py v1 <language(optional)>
+ ```
+Or maunally switch version in WebUI
+
+ ### Finetune
+
+ #### Path Auto-filling is now supported
+
+     1.Fill in the audio path
+
+     2.Slice the audio into small chunks
+
+     3.Denoise(optinal)
+
+     4.ASR
+
+     5.Proofreading ASR transcriptions
+
+     6.Go to the next Tab, then finetune the model
+
+ ### Open Inference WebUI
+
+ #### Integrated Package Users
+
+ Double-click `go-webui-v2.bat` or use `go-webui-v2.ps1` ,then open the inference webui at  `1-GPT-SoVITS-TTS/1C-inference` 
+
+ #### Others
+
+ ```bash
+ python moyoyo_tts/inference_webui.py <language(optional)>
+ ```
+ OR
+
+ ```bash
+ python webui.py
+ ```
+then open the inference webui at `1-GPT-SoVITS-TTS/1C-inference`
+
+ ## V2 Release Notes
+
+New Features:
+
+1. Support Korean and Cantonese
+
+2. An optimized text frontend
+
+3. Pre-trained model extended from 2k hours to 5k hours
+
+4. Improved synthesis quality for low-quality reference audio 
+
+    [more details](https://github.com/RVC-Boss/GPT-SoVITS/wiki/GPT%E2%80%90SoVITS%E2%80%90v2%E2%80%90features-(%E6%96%B0%E7%89%B9%E6%80%A7) ) 
+
+Use v2 from v1 environment: 
+
+1. `pip install -r requirements.txt` to update some packages
+
+2. Clone the latest codes from github.
+
+3. Download v2 pretrained models from [huggingface](https://huggingface.co/lj1995/GPT-SoVITS/tree/main/gsv-v2final-pretrained) and put them into `GPT_SoVITS\pretrained_models\gsv-v2final-pretrained`.
+
+    Chinese v2 additional: [G2PWModel_1.1.zip](https://paddlespeech.bj.bcebos.com/Parakeet/released_models/g2p/G2PWModel_1.1.zip)（Download G2PW models,  unzip and rename to `G2PWModel`, and then place them in `GPT_SoVITS/text`.
+     
+## Todo List
+
+- [x] **High Priority:**
+
+  - [x] Localization in Japanese and English.
+  - [x] User guide.
+  - [x] Japanese and English dataset fine tune training.
+
+- [ ] **Features:**
+  - [x] Zero-shot voice conversion (5s) / few-shot voice conversion (1min).
+  - [x] TTS speaking speed control.
+  - [ ] ~~Enhanced TTS emotion control.~~
+  - [ ] Experiment with changing SoVITS token inputs to probability distribution of GPT vocabs (transformer latent).
+  - [x] Improve English and Japanese text frontend.
+  - [ ] Develop tiny and larger-sized TTS models.
+  - [x] Colab scripts.
+  - [ ] Try expand training dataset (2k hours -> 10k hours).
+  - [x] better sovits base model (enhanced audio quality)
+  - [ ] model mix
+
+## (Additional) Method for running from the command line
+Use the command line to open the WebUI for UVR5
+```
+python tools/uvr5/webui.py "<infer_device>" <is_half> <webui_port_uvr5>
+```
+<!-- If you can't open a browser, follow the format below for UVR processing,This is using mdxnet for audio processing
+```
+python mdxnet.py --model --input_root --output_vocal --output_ins --agg_level --format --device --is_half_precision 
+``` -->
+This is how the audio segmentation of the dataset is done using the command line
+```
+python audio_slicer.py \
+    --input_path "<path_to_original_audio_file_or_directory>" \
+    --output_root "<directory_where_subdivided_audio_clips_will_be_saved>" \
+    --threshold <volume_threshold> \
+    --min_length <minimum_duration_of_each_subclip> \
+    --min_interval <shortest_time_gap_between_adjacent_subclips> 
+    --hop_size <step_size_for_computing_volume_curve>
+```
+This is how dataset ASR processing is done using the command line(Only Chinese)
+```
+python tools/asr/funasr_asr.py -i <input> -o <output>
+```
+ASR processing is performed through Faster_Whisper(ASR marking except Chinese)
+
+(No progress bars, GPU performance may cause time delays)
+```
+python ./tools/asr/fasterwhisper_asr.py -i <input> -o <output> -l <language> -p <precision>
+```
+A custom list save path is enabled
+
+## Credits
+
+Special thanks to the following projects and contributors:
+
+### Theoretical Research
+- [ar-vits](https://github.com/innnky/ar-vits)
+- [SoundStorm](https://github.com/yangdongchao/SoundStorm/tree/master/soundstorm/s1/AR)
+- [vits](https://github.com/jaywalnut310/vits)
+- [TransferTTS](https://github.com/hcy71o/TransferTTS/blob/master/models.py#L556)
+- [contentvec](https://github.com/auspicious3000/contentvec/)
+- [hifi-gan](https://github.com/jik876/hifi-gan)
+- [fish-speech](https://github.com/fishaudio/fish-speech/blob/main/tools/llama/generate.py#L41)
+### Pretrained Models
+- [Chinese Speech Pretrain](https://github.com/TencentGameMate/chinese_speech_pretrain)
+- [Chinese-Roberta-WWM-Ext-Large](https://huggingface.co/hfl/chinese-roberta-wwm-ext-large)
+### Text Frontend for Inference
+- [paddlespeech zh_normalization](https://github.com/PaddlePaddle/PaddleSpeech/tree/develop/paddlespeech/t2s/frontend/zh_normalization)
+- [LangSegment](https://github.com/juntaosun/LangSegment)
+- [g2pW](https://github.com/GitYCC/g2pW)
+- [pypinyin-g2pW](https://github.com/mozillazg/pypinyin-g2pW)
+- [paddlespeech g2pw](https://github.com/PaddlePaddle/PaddleSpeech/tree/develop/paddlespeech/t2s/frontend/g2pw)
+### WebUI Tools
+- [ultimatevocalremovergui](https://github.com/Anjok07/ultimatevocalremovergui)
+- [audio-slicer](https://github.com/openvpi/audio-slicer)
+- [SubFix](https://github.com/cronrpc/SubFix)
+- [FFmpeg](https://github.com/FFmpeg/FFmpeg)
+- [gradio](https://github.com/gradio-app/gradio)
+- [faster-whisper](https://github.com/SYSTRAN/faster-whisper)
+- [FunASR](https://github.com/alibaba-damo-academy/FunASR)
+
+Thankful to @Naozumi520 for providing the Cantonese training set and for the guidance on Cantonese-related knowledge.
+
+## Thanks to all contributors for their efforts
+
+<a href="https://github.com/RVC-Boss/GPT-SoVITS/graphs/contributors" target="_blank">
+  <img src="https://contrib.rocks/image?repo=RVC-Boss/GPT-SoVITS" />
+</a>

TTS_infer_pack/TTS.py

@@ -0,0 +1,1047 @@
+import gc
+import math
+import os
+import random
+import sys
+import traceback
+from copy import deepcopy
+from time import time as ttime
+from typing import List, Tuple, Union
+
+import ffmpeg
+import librosa
+import soundfile
+import numpy as np
+import torch
+import yaml
+from tqdm import tqdm
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+
+from moyoyo_tts.AR.models.t2s_lightning_module import Text2SemanticLightningModule
+from moyoyo_tts.TTS_infer_pack.TextPreprocessor import TextPreprocessor
+from moyoyo_tts.TTS_infer_pack.text_segmentation_method import splits
+from moyoyo_tts.feature_extractor.cnhubert import CNHubert
+from moyoyo_tts.module.mel_processing import spectrogram_torch
+from moyoyo_tts.module.models import SynthesizerTrn
+from moyoyo_tts.tools.i18n.i18n import I18nAuto, scan_language_list
+from moyoyo_tts.tools.my_utils import load_audio
+
+language = os.environ.get("language", "Auto")
+language = sys.argv[-1] if sys.argv[-1] in scan_language_list() else language
+i18n = I18nAuto(language=language)
+
+# configs/tts_infer.yaml
+"""
+custom:
+  bert_base_path: moyoyo_tts/pretrained_models/chinese-roberta-wwm-ext-large
+  cnhuhbert_base_path: moyoyo_tts/pretrained_models/chinese-hubert-base
+  device: cpu
+  is_half: false
+  t2s_weights_path: moyoyo_tts/pretrained_models/gsv-v2final-pretrained/s1bert25hz-5kh-longer-epoch=12-step=369668.ckpt
+  vits_weights_path: moyoyo_tts/pretrained_models/gsv-v2final-pretrained/s2G2333k.pth
+  version: v2
+default:
+  bert_base_path: moyoyo_tts/pretrained_models/chinese-roberta-wwm-ext-large
+  cnhuhbert_base_path: moyoyo_tts/pretrained_models/chinese-hubert-base
+  device: cpu
+  is_half: false
+  t2s_weights_path: moyoyo_tts/pretrained_models/s1bert25hz-2kh-longer-epoch=68e-step=50232.ckpt
+  vits_weights_path: moyoyo_tts/pretrained_models/s2G488k.pth
+  version: v1
+default_v2:
+  bert_base_path: moyoyo_tts/pretrained_models/chinese-roberta-wwm-ext-large
+  cnhuhbert_base_path: moyoyo_tts/pretrained_models/chinese-hubert-base
+  device: cpu
+  is_half: false
+  t2s_weights_path: moyoyo_tts/pretrained_models/gsv-v2final-pretrained/s1bert25hz-5kh-longer-epoch=12-step=369668.ckpt
+  vits_weights_path: moyoyo_tts/pretrained_models/gsv-v2final-pretrained/s2G2333k.pth
+  version: v2
+"""
+
+
+def set_seed(seed: int):
+    seed = int(seed)
+    seed = seed if seed != -1 else random.randrange(1 << 32)
+    #print(f"Set seed to {seed}")
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    try:
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed(seed)
+            torch.cuda.manual_seed_all(seed)
+            # torch.backends.cudnn.deterministic = True
+            # torch.backends.cudnn.benchmark = False
+            # torch.backends.cudnn.enabled = True
+            # 开启后会影响精度
+            torch.backends.cuda.matmul.allow_tf32 = False
+            torch.backends.cudnn.allow_tf32 = False
+    except:
+        pass
+    return seed
+
+
+class TTS_Config:
+    default_configs = {
+        "default": {
+            "device": "cpu",
+            "is_half": False,
+            "version": "v1",
+            "t2s_weights_path": "moyoyo_tts/pretrained_models/s1bert25hz-2kh-longer-epoch=68e-step=50232.ckpt",
+            "vits_weights_path": "moyoyo_tts/pretrained_models/s2G488k.pth",
+            "cnhuhbert_base_path": "moyoyo_tts/pretrained_models/chinese-hubert-base",
+            "bert_base_path": "moyoyo_tts/pretrained_models/chinese-roberta-wwm-ext-large",
+        },
+        "default_v2": {
+            "device": "cpu",
+            "is_half": False,
+            "version": "v2",
+            "t2s_weights_path": "moyoyo_tts/pretrained_models/gsv-v2final-pretrained/s1bert25hz-5kh-longer-epoch=12-step=369668.ckpt",
+            "vits_weights_path": "moyoyo_tts/pretrained_models/gsv-v2final-pretrained/s2G2333k.pth",
+            "cnhuhbert_base_path": "moyoyo_tts/pretrained_models/chinese-hubert-base",
+            "bert_base_path": "moyoyo_tts/pretrained_models/chinese-roberta-wwm-ext-large",
+        },
+    }
+    configs: dict = None
+    v1_languages: list = ["auto", "en", "zh", "ja", "all_zh", "all_ja"]
+    v2_languages: list = ["auto", "auto_yue", "en", "zh", "ja", "yue", "ko", "all_zh", "all_ja", "all_yue", "all_ko"]
+    languages: list = v2_languages
+
+    # "all_zh",#全部按中文识别
+    # "en",#全部按英文识别#######不变
+    # "all_ja",#全部按日文识别
+    # "all_yue",#全部按中文识别
+    # "all_ko",#全部按韩文识别
+    # "zh",#按中英混合识别####不变
+    # "ja",#按日英混合识别####不变
+    # "yue",#按粤英混合识别####不变
+    # "ko",#按韩英混合识别####不变
+    # "auto",#多语种启动切分识别语种
+    # "auto_yue",#多语种启动切分识别语种
+
+    def __init__(self, configs: Union[dict, str] = None):
+
+        if configs in ["", None]:
+            # 设置默认配置文件路径
+            configs_base_path: str = "moyoyo_tts/configs/"
+            os.makedirs(configs_base_path, exist_ok=True)
+            self.configs_path: str = os.path.join(configs_base_path, "tts_infer.yaml")
+            if not os.path.exists(self.configs_path):
+                self.save_configs()
+                print(f"Create default config file at {self.configs_path}")
+            configs: dict = deepcopy(self.default_configs)
+
+        if isinstance(configs, str):
+            self.configs_path = configs
+            configs: dict = self._load_configs(self.configs_path)
+
+        assert isinstance(configs, dict)
+        version = configs.get("version", "v2").lower()
+        assert version in ["v1", "v2"]
+        self.default_configs["default"] = configs.get("default", self.default_configs["default"])
+        self.default_configs["default_v2"] = configs.get("default_v2", self.default_configs["default_v2"])
+
+        default_config_key = "default" if version == "v1" else "default_v2"
+        self.configs: dict = configs.get("custom", deepcopy(self.default_configs[default_config_key]))
+
+        self.device = self.configs.get("device", torch.device("cpu"))
+        self.is_half = self.configs.get("is_half", False)
+        self.version = version
+        self.t2s_weights_path = self.configs.get("t2s_weights_path", None)
+        self.vits_weights_path = self.configs.get("vits_weights_path", None)
+        # self.prompt_semantic_path = self.configs.get('prompt_semantic_path')
+        # self.refer_spec_path = self.configs.get('refer_spec_path')
+        self.bert_base_path = self.configs.get("bert_base_path", None)
+        os.environ['bert_path'] = f'{self.bert_base_path}'
+        self.cnhuhbert_base_path = self.configs.get("cnhuhbert_base_path", None)
+        self.languages = self.v2_languages if self.version == "v2" else self.v1_languages
+
+        if (self.t2s_weights_path in [None, ""]) or (not os.path.exists(self.t2s_weights_path)):
+            self.t2s_weights_path = self.default_configs[default_config_key]['t2s_weights_path']
+            print(f"fall back to default t2s_weights_path: {self.t2s_weights_path}")
+
+        if (self.vits_weights_path in [None, ""]) or (not os.path.exists(self.vits_weights_path)):
+            self.vits_weights_path = self.default_configs[default_config_key]['vits_weights_path']
+            print(f"fall back to default vits_weights_path: {self.vits_weights_path}")
+
+        if (self.bert_base_path in [None, ""]) or (not os.path.exists(self.bert_base_path)):
+            self.bert_base_path = self.default_configs[default_config_key]['bert_base_path']
+            print(f"fall back to default bert_base_path: {self.bert_base_path}")
+
+        if (self.cnhuhbert_base_path in [None, ""]) or (not os.path.exists(self.cnhuhbert_base_path)):
+            self.cnhuhbert_base_path = self.default_configs[default_config_key]['cnhuhbert_base_path']
+            print(f"fall back to default cnhuhbert_base_path: {self.cnhuhbert_base_path}")
+
+        # if (self.prompt_semantic_path in [None, ""]) or (not os.path.exists(self.prompt_semantic_path)):
+        #     self.prompt_semantic_path = self.default_configs['prompt_semantic_path']
+        #     print(f"fall back to default prompt_semantic_path: {self.prompt_semantic_path}")
+        #
+        # if (self.refer_spec_path in [None, ""]) or (not os.path.exists(self.refer_spec_path)):
+        #     self.refer_spec_path = self.default_configs['refer_spec_path']
+        #     print(f"fall back to default refer_spec_path: {self.refer_spec_path}")
+
+        self.update_configs()
+
+        self.max_sec = None
+        self.hz: int = 50
+        self.semantic_frame_rate: str = "25hz"
+        self.segment_size: int = 20480
+        self.filter_length: int = 2048
+        self.sampling_rate: int = 32000
+        self.hop_length: int = 640
+        self.win_length: int = 2048
+        self.n_speakers: int = 300
+
+    def _load_configs(self, configs_path: str) -> dict:
+        if os.path.exists(configs_path):
+            ...
+        else:
+            print(i18n("路径不存在,使用默认配置"))
+            self.save_configs(configs_path)
+        with open(configs_path, 'r') as f:
+            configs = yaml.load(f, Loader=yaml.FullLoader)
+
+        return configs
+
+    def save_configs(self, configs_path: str = None) -> None:
+        configs = deepcopy(self.default_configs)
+        if self.configs is not None:
+            configs["custom"] = self.update_configs()
+
+        if configs_path is None and hasattr(self, 'configs_path'):
+            configs_path = self.configs_path
+
+        if configs_path is None:
+            return
+
+        with open(configs_path, 'w') as f:
+            yaml.dump(configs, f)
+
+    def update_configs(self):
+        self.config = {
+            "device": str(self.device),
+            "is_half": self.is_half,
+            "version": self.version,
+            "t2s_weights_path": self.t2s_weights_path,
+            "vits_weights_path": self.vits_weights_path,
+            "bert_base_path": self.bert_base_path,
+            "cnhuhbert_base_path": self.cnhuhbert_base_path,
+        }
+        return self.config
+
+    def update_version(self, version: str) -> None:
+        self.version = version
+        self.languages = self.v2_languages if self.version == "v2" else self.v1_languages
+
+    def __str__(self):
+        self.configs = self.update_configs()
+        string = "TTS Config".center(100, '-') + '\n'
+        for k, v in self.configs.items():
+            string += f"{str(k).ljust(20)}: {str(v)}\n"
+        string += "-" * 100 + '\n'
+        return string
+
+    def __repr__(self):
+        return self.__str__()
+
+    def __hash__(self):
+        return hash(self.configs_path)
+
+    def __eq__(self, other):
+        return isinstance(other, TTS_Config) and self.configs_path == other.configs_path
+
+
+class TTS:
+    def __init__(self, configs: Union[dict, str, TTS_Config]):
+        if isinstance(configs, TTS_Config):
+            self.configs = configs
+        else:
+            self.configs: TTS_Config = TTS_Config(configs)
+
+        self.t2s_model: Text2SemanticLightningModule = None
+        self.vits_model: SynthesizerTrn = None
+        self.bert_tokenizer: AutoTokenizer = None
+        self.bert_model: AutoModelForMaskedLM = None
+        self.cnhuhbert_model: CNHubert = None
+
+        self._init_models()
+
+        self.text_preprocessor: TextPreprocessor = \
+            TextPreprocessor(self.bert_model,
+                             self.bert_tokenizer,
+                             self.configs.device)
+
+        self.prompt_cache: dict = {
+            "ref_audio_path": None,
+            "prompt_semantic": None,
+            "refer_spec": [],
+            "prompt_text": None,
+            "prompt_lang": None,
+            "phones": None,
+            "bert_features": None,
+            "norm_text": None,
+            "aux_ref_audio_paths": [],
+        }
+
+        self.stop_flag: bool = False
+        self.precision: torch.dtype = torch.float16 if self.configs.is_half else torch.float32
+
+    def _init_models(self, ):
+        self.init_t2s_weights(self.configs.t2s_weights_path)
+        self.init_vits_weights(self.configs.vits_weights_path)
+        self.init_bert_weights(self.configs.bert_base_path)
+        self.init_cnhuhbert_weights(self.configs.cnhuhbert_base_path)
+        # self.enable_half_precision(self.configs.is_half)
+
+    def init_cnhuhbert_weights(self, base_path: str):
+        print(f"Loading CNHuBERT weights from {base_path}")
+        self.cnhuhbert_model = CNHubert(base_path)
+        self.cnhuhbert_model = self.cnhuhbert_model.eval()
+        self.cnhuhbert_model = self.cnhuhbert_model.to(self.configs.device)
+        if self.configs.is_half and str(self.configs.device) != "cpu":
+            self.cnhuhbert_model = self.cnhuhbert_model.half()
+
+    def init_bert_weights(self, base_path: str):
+        print(f"Loading BERT weights from {base_path}")
+        self.bert_tokenizer = AutoTokenizer.from_pretrained(base_path)
+        self.bert_model = AutoModelForMaskedLM.from_pretrained(base_path)
+        self.bert_model = self.bert_model.eval()
+        self.bert_model = self.bert_model.to(self.configs.device)
+        if self.configs.is_half and str(self.configs.device) != "cpu":
+            self.bert_model = self.bert_model.half()
+
+    def init_vits_weights(self, weights_path: str):
+        print(f"Loading VITS weights from {weights_path}")
+        self.configs.vits_weights_path = weights_path
+        dict_s2 = torch.load(weights_path, map_location=self.configs.device)
+        hps = dict_s2["config"]
+        if dict_s2['weight']['enc_p.text_embedding.weight'].shape[0] == 322:
+            self.configs.update_version("v1")
+        else:
+            self.configs.update_version("v2")
+        self.configs.save_configs()
+
+        hps["model"]["version"] = self.configs.version
+        self.configs.filter_length = hps["data"]["filter_length"]
+        self.configs.segment_size = hps["train"]["segment_size"]
+        self.configs.sampling_rate = hps["data"]["sampling_rate"]
+        self.configs.hop_length = hps["data"]["hop_length"]
+        self.configs.win_length = hps["data"]["win_length"]
+        self.configs.n_speakers = hps["data"]["n_speakers"]
+        self.configs.semantic_frame_rate = "25hz"
+        kwargs = hps["model"]
+        vits_model = SynthesizerTrn(
+            self.configs.filter_length // 2 + 1,
+            self.configs.segment_size // self.configs.hop_length,
+            n_speakers=self.configs.n_speakers,
+            **kwargs
+        )
+
+        if hasattr(vits_model, "enc_q"):
+            del vits_model.enc_q
+
+        vits_model = vits_model.to(self.configs.device)
+        vits_model = vits_model.eval()
+        vits_model.load_state_dict(dict_s2["weight"], strict=False)
+        self.vits_model = vits_model
+        if self.configs.is_half and str(self.configs.device) != "cpu":
+            self.vits_model = self.vits_model.half()
+
+    def init_t2s_weights(self, weights_path: str):
+        print(f"Loading Text2Semantic weights from {weights_path}")
+        self.configs.t2s_weights_path = weights_path
+        self.configs.save_configs()
+        self.configs.hz = 50
+        dict_s1 = torch.load(weights_path, map_location=self.configs.device)
+        config = dict_s1["config"]
+        self.configs.max_sec = config["data"]["max_sec"]
+        t2s_model = Text2SemanticLightningModule(config, "****", is_train=False)
+        t2s_model.load_state_dict(dict_s1["weight"])
+        t2s_model = t2s_model.to(self.configs.device)
+        t2s_model = t2s_model.eval()
+        self.t2s_model = t2s_model
+        if self.configs.is_half and str(self.configs.device) != "cpu":
+            self.t2s_model = self.t2s_model.half()
+
+    def enable_half_precision(self, enable: bool = True, save: bool = True):
+        '''
+            To enable half precision for the TTS model.
+            Args:
+                enable: bool, whether to enable half precision.
+
+        '''
+        if str(self.configs.device) == "cpu" and enable:
+            print("Half precision is not supported on CPU.")
+            return
+
+        self.configs.is_half = enable
+        self.precision = torch.float16 if enable else torch.float32
+        if save:
+            self.configs.save_configs()
+        if enable:
+            if self.t2s_model is not None:
+                self.t2s_model = self.t2s_model.half()
+            if self.vits_model is not None:
+                self.vits_model = self.vits_model.half()
+            if self.bert_model is not None:
+                self.bert_model = self.bert_model.half()
+            if self.cnhuhbert_model is not None:
+                self.cnhuhbert_model = self.cnhuhbert_model.half()
+        else:
+            if self.t2s_model is not None:
+                self.t2s_model = self.t2s_model.float()
+            if self.vits_model is not None:
+                self.vits_model = self.vits_model.float()
+            if self.bert_model is not None:
+                self.bert_model = self.bert_model.float()
+            if self.cnhuhbert_model is not None:
+                self.cnhuhbert_model = self.cnhuhbert_model.float()
+
+    def set_device(self, device: torch.device, save: bool = True):
+        '''
+            To set the device for all models.
+            Args:
+                device: torch.device, the device to use for all models.
+        '''
+        self.configs.device = device
+        if save:
+            self.configs.save_configs()
+        if self.t2s_model is not None:
+            self.t2s_model = self.t2s_model.to(device)
+        if self.vits_model is not None:
+            self.vits_model = self.vits_model.to(device)
+        if self.bert_model is not None:
+            self.bert_model = self.bert_model.to(device)
+        if self.cnhuhbert_model is not None:
+            self.cnhuhbert_model = self.cnhuhbert_model.to(device)
+
+    def set_ref_audio(self, ref_audio_path: str):
+        '''
+            To set the reference audio for the TTS model,
+                including the prompt_semantic and refer_spepc.
+            Args:
+                ref_audio_path: str, the path of the reference audio.
+        '''
+        self._set_prompt_semantic(ref_audio_path)
+        self._set_ref_spec(ref_audio_path)
+        self._set_ref_audio_path(ref_audio_path)
+        # self._set_ref_audio_path(ref_audio_path)
+        # self.prompt_cache['prompt_semantic'] = torch.load(self.configs.prompt_semantic_path, map_location=self.configs.device)
+        # self.prompt_cache["refer_spec"] = [torch.load(self.configs.refer_spec_path, map_location=self.configs.device)]
+
+    def _set_ref_audio_path(self, ref_audio_path):
+        self.prompt_cache["ref_audio_path"] = ref_audio_path
+
+    def _set_ref_spec(self, ref_audio_path):
+        spec = self._get_ref_spec(ref_audio_path)
+        if self.prompt_cache["refer_spec"] in [[], None]:
+            self.prompt_cache["refer_spec"] = [spec]
+        else:
+            self.prompt_cache["refer_spec"][0] = spec
+
+    def _get_ref_spec(self, ref_audio_path):
+        audio = load_audio(ref_audio_path, int(self.configs.sampling_rate))
+        audio = torch.FloatTensor(audio)
+        maxx = audio.abs().max()
+        if (maxx > 1): audio /= min(2, maxx)
+        audio_norm = audio
+        audio_norm = audio_norm.unsqueeze(0)
+        spec = spectrogram_torch(
+            audio_norm,
+            self.configs.filter_length,
+            self.configs.sampling_rate,
+            self.configs.hop_length,
+            self.configs.win_length,
+            center=False,
+        )
+        spec = spec.to(self.configs.device)
+        if self.configs.is_half:
+            spec = spec.half()
+        return spec
+
+    def _set_prompt_semantic(self, ref_wav_path: str):
+        zero_wav = np.zeros(
+            int(self.configs.sampling_rate * 0.3),
+            dtype=np.float16 if self.configs.is_half else np.float32,
+        )
+        with torch.no_grad():
+            wav16k, sr = librosa.load(ref_wav_path, sr=16000)
+            y, sr, = soundfile.read(ref_wav_path)
+            duration = librosa.get_duration(y=y, sr=sr)
+            if (duration < 3 or duration > 10):
+                raise OSError(i18n("参考音频在3~10秒范围外，请更换！"))
+            wav16k = torch.from_numpy(wav16k)
+            zero_wav_torch = torch.from_numpy(zero_wav)
+            wav16k = wav16k.to(self.configs.device)
+            zero_wav_torch = zero_wav_torch.to(self.configs.device)
+            if self.configs.is_half:
+                wav16k = wav16k.half()
+                zero_wav_torch = zero_wav_torch.half()
+
+            wav16k = torch.cat([wav16k, zero_wav_torch])
+            hubert_feature = self.cnhuhbert_model.model(wav16k.unsqueeze(0))[
+                "last_hidden_state"
+            ].transpose(
+                1, 2
+            )  # .float()
+            codes = self.vits_model.extract_latent(hubert_feature)
+
+            prompt_semantic = codes[0, 0].to(self.configs.device)
+            self.prompt_cache["prompt_semantic"] = prompt_semantic
+
+    def batch_sequences(self, sequences: List[torch.Tensor], axis: int = 0, pad_value: int = 0, max_length: int = None):
+        seq = sequences[0]
+        ndim = seq.dim()
+        if axis < 0:
+            axis += ndim
+        dtype: torch.dtype = seq.dtype
+        pad_value = torch.tensor(pad_value, dtype=dtype)
+        seq_lengths = [seq.shape[axis] for seq in sequences]
+        if max_length is None:
+            max_length = max(seq_lengths)
+        else:
+            max_length = max(seq_lengths) if max_length < max(seq_lengths) else max_length
+
+        padded_sequences = []
+        for seq, length in zip(sequences, seq_lengths):
+            padding = [0] * axis + [0, max_length - length] + [0] * (ndim - axis - 1)
+            padded_seq = torch.nn.functional.pad(seq, padding, value=pad_value)
+            padded_sequences.append(padded_seq)
+        batch = torch.stack(padded_sequences)
+        return batch
+
+    def to_batch(self, data: list,
+                 prompt_data: dict = None,
+                 batch_size: int = 5,
+                 threshold: float = 0.75,
+                 split_bucket: bool = True,
+                 device: torch.device = torch.device("cpu"),
+                 precision: torch.dtype = torch.float32,
+                 ):
+        _data: list = []
+        index_and_len_list = []
+        for idx, item in enumerate(data):
+            norm_text_len = len(item["norm_text"])
+            index_and_len_list.append([idx, norm_text_len])
+
+        batch_index_list = []
+        if split_bucket:
+            index_and_len_list.sort(key=lambda x: x[1])
+            index_and_len_list = np.array(index_and_len_list, dtype=np.int64)
+
+            batch_index_list_len = 0
+            pos = 0
+            while pos < index_and_len_list.shape[0]:
+                # batch_index_list.append(index_and_len_list[pos:min(pos+batch_size,len(index_and_len_list))])
+                pos_end = min(pos + batch_size, index_and_len_list.shape[0])
+                while pos < pos_end:
+                    batch = index_and_len_list[pos:pos_end, 1].astype(np.float32)
+                    score = batch[(pos_end - pos) // 2] / (batch.mean() + 1e-8)
+                    if (score >= threshold) or (pos_end - pos == 1):
+                        batch_index = index_and_len_list[pos:pos_end, 0].tolist()
+                        batch_index_list_len += len(batch_index)
+                        batch_index_list.append(batch_index)
+                        pos = pos_end
+                        break
+                    pos_end = pos_end - 1
+
+            assert batch_index_list_len == len(data)
+
+        else:
+            for i in range(len(data)):
+                if i % batch_size == 0:
+                    batch_index_list.append([])
+                batch_index_list[-1].append(i)
+
+        for batch_idx, index_list in enumerate(batch_index_list):
+            item_list = [data[idx] for idx in index_list]
+            phones_list = []
+            phones_len_list = []
+            # bert_features_list = []
+            all_phones_list = []
+            all_phones_len_list = []
+            all_bert_features_list = []
+            norm_text_batch = []
+            all_bert_max_len = 0
+            all_phones_max_len = 0
+            for item in item_list:
+                if prompt_data is not None:
+                    all_bert_features = torch.cat([prompt_data["bert_features"], item["bert_features"]], 1) \
+                        .to(dtype=precision, device=device)
+                    all_phones = torch.LongTensor(prompt_data["phones"] + item["phones"]).to(device)
+                    phones = torch.LongTensor(item["phones"]).to(device)
+                    # norm_text = prompt_data["norm_text"]+item["norm_text"]
+                else:
+                    all_bert_features = item["bert_features"] \
+                        .to(dtype=precision, device=device)
+                    phones = torch.LongTensor(item["phones"]).to(device)
+                    all_phones = phones
+                    # norm_text = item["norm_text"]
+
+                all_bert_max_len = max(all_bert_max_len, all_bert_features.shape[-1])
+                all_phones_max_len = max(all_phones_max_len, all_phones.shape[-1])
+
+                phones_list.append(phones)
+                phones_len_list.append(phones.shape[-1])
+                all_phones_list.append(all_phones)
+                all_phones_len_list.append(all_phones.shape[-1])
+                all_bert_features_list.append(all_bert_features)
+                norm_text_batch.append(item["norm_text"])
+
+            phones_batch = phones_list
+            all_phones_batch = all_phones_list
+            all_bert_features_batch = all_bert_features_list
+
+            max_len = max(all_bert_max_len, all_phones_max_len)
+            # phones_batch = self.batch_sequences(phones_list, axis=0, pad_value=0, max_length=max_len)
+            #### 直接对phones和bert_features进行pad。（padding策略会影响T2S模型生成的结果，但不直接影响复读概率。影响复读概率的主要因素是mask的策略）
+            # all_phones_batch = self.batch_sequences(all_phones_list, axis=0, pad_value=0, max_length=max_len)
+            # all_bert_features_batch = all_bert_features_list
+            # all_bert_features_batch = torch.zeros((len(all_bert_features_list), 1024, max_len), dtype=precision, device=device)
+            # for idx, item in enumerate(all_bert_features_list):
+            #     all_bert_features_batch[idx, :, : item.shape[-1]] = item
+
+            # #### 先对phones进行embedding、对bert_features进行project，再pad到相同长度，（padding策略会影响T2S模型生成的结果，但不直接影响复读概率。影响复读概率的主要因素是mask的策略）
+            # all_phones_list = [self.t2s_model.model.ar_text_embedding(item.to(self.t2s_model.device)) for item in all_phones_list]
+            # all_phones_list = [F.pad(item,(0,0,0,max_len-item.shape[0]),value=0) for item in all_phones_list]
+            # all_phones_batch = torch.stack(all_phones_list, dim=0)
+
+            # all_bert_features_list = [self.t2s_model.model.bert_proj(item.to(self.t2s_model.device).transpose(0, 1)) for item in all_bert_features_list]
+            # all_bert_features_list = [F.pad(item,(0,0,0,max_len-item.shape[0]), value=0) for item in all_bert_features_list]
+            # all_bert_features_batch = torch.stack(all_bert_features_list, dim=0)
+
+            batch = {
+                "phones": phones_batch,
+                "phones_len": torch.LongTensor(phones_len_list).to(device),
+                "all_phones": all_phones_batch,
+                "all_phones_len": torch.LongTensor(all_phones_len_list).to(device),
+                "all_bert_features": all_bert_features_batch,
+                "norm_text": norm_text_batch,
+                "max_len": max_len,
+            }
+            _data.append(batch)
+
+        return _data, batch_index_list
+
+    def recovery_order(self, data: list, batch_index_list: list) -> list:
+        '''
+        Recovery the order of the audio according to the batch_index_list.
+
+        Args:
+            data (List[list(np.ndarray)]): the out of order audio .
+            batch_index_list (List[list[int]]): the batch index list.
+
+        Returns:
+            list (List[np.ndarray]): the data in the original order.
+        '''
+        length = len(sum(batch_index_list, []))
+        _data = [None] * length
+        for i, index_list in enumerate(batch_index_list):
+            for j, index in enumerate(index_list):
+                _data[index] = data[i][j]
+        return _data
+
+    def stop(self, ):
+        '''
+        Stop the inference process.
+        '''
+        self.stop_flag = True
+
+    @torch.no_grad()
+    def run(self, inputs: dict):
+        """
+        Text to speech inference.
+
+        Args:
+            inputs (dict):
+                {
+                    "text": "",                   # str.(required) text to be synthesized
+                    "text_lang: "",               # str.(required) language of the text to be synthesized
+                    "ref_audio_path": "",         # str.(required) reference audio path
+                    "aux_ref_audio_paths": [],    # list.(optional) auxiliary reference audio paths for multi-speaker tone fusion
+                    "prompt_text": "",            # str.(optional) prompt text for the reference audio
+                    "prompt_lang": "",            # str.(required) language of the prompt text for the reference audio
+                    "top_k": 5,                   # int. top k sampling
+                    "top_p": 1,                   # float. top p sampling
+                    "temperature": 1,             # float. temperature for sampling
+                    "text_split_method": "cut0",  # str. text split method, see text_segmentation_method.py for details.
+                    "batch_size": 1,              # int. batch size for inference
+                    "batch_threshold": 0.75,      # float. threshold for batch splitting.
+                    "split_bucket: True,          # bool. whether to split the batch into multiple buckets.
+                    "return_fragment": False,     # bool. step by step return the audio fragment.
+                    "speed_factor":1.0,           # float. control the speed of the synthesized audio.
+                    "fragment_interval":0.3,      # float. to control the interval of the audio fragment.
+                    "seed": -1,                   # int. random seed for reproducibility.
+                    "parallel_infer": True,       # bool. whether to use parallel inference.
+                    "repetition_penalty": 1.35    # float. repetition penalty for T2S model.
+                }
+        returns:
+            Tuple[int, np.ndarray]: sampling rate and audio data.
+        """
+        ########## variables initialization ###########
+        self.stop_flag: bool = False
+        text: str = inputs.get("text", "")
+        text_lang: str = inputs.get("text_lang", "")
+        ref_audio_path: str = inputs.get("ref_audio_path", "")
+        aux_ref_audio_paths: list = inputs.get("aux_ref_audio_paths", [])
+        prompt_text: str = inputs.get("prompt_text", "")
+        prompt_lang: str = inputs.get("prompt_lang", "")
+        top_k: int = inputs.get("top_k", 5)
+        top_p: float = inputs.get("top_p", 1)
+        temperature: float = inputs.get("temperature", 1)
+        text_split_method: str = inputs.get("text_split_method", "cut0")
+        batch_size = inputs.get("batch_size", 1)
+        batch_threshold = inputs.get("batch_threshold", 0.75)
+        speed_factor = inputs.get("speed_factor", 1.0)
+        split_bucket = inputs.get("split_bucket", True)
+        return_fragment = inputs.get("return_fragment", False)
+        fragment_interval = inputs.get("fragment_interval", 0.3)
+        seed = inputs.get("seed", -1)
+        seed = -1 if seed in ["", None] else seed
+        actual_seed = set_seed(seed)
+        parallel_infer = inputs.get("parallel_infer", True)
+        repetition_penalty = inputs.get("repetition_penalty", 1.35)
+
+        if parallel_infer:
+            #print(i18n("并行推理模式已开启"))
+            self.t2s_model.model.infer_panel = self.t2s_model.model.infer_panel_batch_infer
+        else:
+            #print(i18n("并行推理模式已关闭"))
+            self.t2s_model.model.infer_panel = self.t2s_model.model.infer_panel_naive_batched
+
+        if return_fragment:
+            #print(i18n("分段返回模式已开启"))
+            if split_bucket:
+                split_bucket = False
+                #print(i18n("分段返回模式不支持分桶处理，已自动关闭分桶处理"))
+
+        if split_bucket and speed_factor == 1.0:
+            print(i18n("分桶处理模式已开启"))
+        elif speed_factor != 1.0:
+            #print(i18n("语速调节不支持分桶处理，已自动关闭分桶处理"))
+            split_bucket = False
+        else:
+            print(i18n("分桶处理模式已关闭"))
+
+        if fragment_interval < 0.01:
+            fragment_interval = 0.01
+            #print(i18n("分段间隔过小，已自动设置为0.01"))
+
+        no_prompt_text = False
+        if prompt_text in [None, ""]:
+            no_prompt_text = True
+
+        assert text_lang in self.configs.languages
+        if not no_prompt_text:
+            assert prompt_lang in self.configs.languages
+
+        if ref_audio_path in [None, ""] and \
+                ((self.prompt_cache["prompt_semantic"] is None) or (self.prompt_cache["refer_spec"] in [None, []])):
+            raise ValueError(
+                "ref_audio_path cannot be empty, when the reference audio is not set using set_ref_audio()")
+
+        ###### setting reference audio and prompt text preprocessing ########
+        t0 = ttime()
+        if (ref_audio_path is not None) and (ref_audio_path != self.prompt_cache["ref_audio_path"]):
+            if not os.path.exists(ref_audio_path):
+                raise ValueError(f"{ref_audio_path} not exists")
+            self.set_ref_audio(ref_audio_path)
+
+        aux_ref_audio_paths = aux_ref_audio_paths if aux_ref_audio_paths is not None else []
+        paths = set(aux_ref_audio_paths) & set(self.prompt_cache["aux_ref_audio_paths"])
+        if not (len(list(paths)) == len(aux_ref_audio_paths) == len(self.prompt_cache["aux_ref_audio_paths"])):
+            self.prompt_cache["aux_ref_audio_paths"] = aux_ref_audio_paths
+            self.prompt_cache["refer_spec"] = [self.prompt_cache["refer_spec"][0]]
+            for path in aux_ref_audio_paths:
+                if path in [None, ""]:
+                    continue
+                if not os.path.exists(path):
+                    print(i18n("音频文件不存在，跳过：{}").format(path))
+                    continue
+                self.prompt_cache["refer_spec"].append(self._get_ref_spec(path))
+
+        if not no_prompt_text:
+            prompt_text = prompt_text.strip("\n")
+            if (prompt_text[-1] not in splits): prompt_text += "。" if prompt_lang != "en" else "."
+            #print(i18n("实际输入的参考文本:"), prompt_text)
+            if self.prompt_cache["prompt_text"] != prompt_text:
+                self.prompt_cache["prompt_text"] = prompt_text
+                self.prompt_cache["prompt_lang"] = prompt_lang
+                phones, bert_features, norm_text = \
+                    self.text_preprocessor.segment_and_extract_feature_for_text(
+                        prompt_text,
+                        prompt_lang,
+                        self.configs.version)
+                self.prompt_cache["phones"] = phones
+                self.prompt_cache["bert_features"] = bert_features
+                self.prompt_cache["norm_text"] = norm_text
+
+        ###### text preprocessing ########
+        t1 = ttime()
+        data: list = None
+        if not return_fragment:
+            data = self.text_preprocessor.preprocess(text, text_lang, text_split_method, self.configs.version)
+            if len(data) == 0:
+                yield self.configs.sampling_rate, np.zeros(int(self.configs.sampling_rate),
+                                                           dtype=np.int16)
+                return
+
+            batch_index_list: list = None
+            data, batch_index_list = self.to_batch(data,
+                                                   prompt_data=self.prompt_cache if not no_prompt_text else None,
+                                                   batch_size=batch_size,
+                                                   threshold=batch_threshold,
+                                                   split_bucket=split_bucket,
+                                                   device=self.configs.device,
+                                                   precision=self.precision
+                                                   )
+        else:
+            #print(i18n("############ 切分文本 ############"))
+            texts = self.text_preprocessor.pre_seg_text(text, text_lang, text_split_method)
+            data = []
+            for i in range(len(texts)):
+                if i % batch_size == 0:
+                    data.append([])
+                data[-1].append(texts[i])
+
+            def make_batch(batch_texts):
+                batch_data = []
+                #print(i18n("############ 提取文本Bert特征 ############"))
+                for text in tqdm(batch_texts):
+                    phones, bert_features, norm_text = self.text_preprocessor.segment_and_extract_feature_for_text(text,
+                                                                                                                   text_lang,
+                                                                                                                   self.configs.version)
+                    if phones is None:
+                        continue
+                    res = {
+                        "phones": phones,
+                        "bert_features": bert_features,
+                        "norm_text": norm_text,
+                    }
+                    batch_data.append(res)
+                if len(batch_data) == 0:
+                    return None
+                batch, _ = self.to_batch(batch_data,
+                                         prompt_data=self.prompt_cache if not no_prompt_text else None,
+                                         batch_size=batch_size,
+                                         threshold=batch_threshold,
+                                         split_bucket=False,
+                                         device=self.configs.device,
+                                         precision=self.precision
+                                         )
+                return batch[0]
+
+        t2 = ttime()
+        try:
+            #print("############ 推理 ############")
+            ###### inference ######
+            t_34 = 0.0
+            t_45 = 0.0
+            audio = []
+            for item in data:
+                t3 = ttime()
+                if return_fragment:
+                    item = make_batch(item)
+                    if item is None:
+                        continue
+
+                batch_phones: List[torch.LongTensor] = item["phones"]
+                # batch_phones:torch.LongTensor = item["phones"]
+                batch_phones_len: torch.LongTensor = item["phones_len"]
+                all_phoneme_ids: torch.LongTensor = item["all_phones"]
+                all_phoneme_lens: torch.LongTensor = item["all_phones_len"]
+                all_bert_features: torch.LongTensor = item["all_bert_features"]
+                norm_text: str = item["norm_text"]
+                max_len = item["max_len"]
+
+                #print(i18n("前端处理后的文本(每句):"), norm_text)
+                if no_prompt_text:
+                    prompt = None
+                else:
+                    prompt = self.prompt_cache["prompt_semantic"].expand(len(all_phoneme_ids), -1).to(
+                        self.configs.device)
+
+                pred_semantic_list, idx_list = self.t2s_model.model.infer_panel(
+                    all_phoneme_ids,
+                    all_phoneme_lens,
+                    prompt,
+                    all_bert_features,
+                    # prompt_phone_len=ph_offset,
+                    top_k=top_k,
+                    top_p=top_p,
+                    temperature=temperature,
+                    early_stop_num=self.configs.hz * self.configs.max_sec,
+                    max_len=max_len,
+                    repetition_penalty=repetition_penalty,
+                )
+                t4 = ttime()
+                t_34 += t4 - t3
+
+                refer_audio_spec: torch.Tensor = [item.to(dtype=self.precision, device=self.configs.device) for item in
+                                                  self.prompt_cache["refer_spec"]]
+
+                batch_audio_fragment = []
+
+                # ## vits并行推理 method 1
+                # pred_semantic_list = [item[-idx:] for item, idx in zip(pred_semantic_list, idx_list)]
+                # pred_semantic_len = torch.LongTensor([item.shape[0] for item in pred_semantic_list]).to(self.configs.device)
+                # pred_semantic = self.batch_sequences(pred_semantic_list, axis=0, pad_value=0).unsqueeze(0)
+                # max_len = 0
+                # for i in range(0, len(batch_phones)):
+                #     max_len = max(max_len, batch_phones[i].shape[-1])
+                # batch_phones = self.batch_sequences(batch_phones, axis=0, pad_value=0, max_length=max_len)
+                # batch_phones = batch_phones.to(self.configs.device)
+                # batch_audio_fragment = (self.vits_model.batched_decode(
+                #         pred_semantic, pred_semantic_len, batch_phones, batch_phones_len,refer_audio_spec
+                #     ))
+
+                if speed_factor == 1.0:
+                    # ## vits并行推理 method 2
+                    pred_semantic_list = [item[-idx:] for item, idx in zip(pred_semantic_list, idx_list)]
+                    upsample_rate = math.prod(self.vits_model.upsample_rates)
+                    audio_frag_idx = [pred_semantic_list[i].shape[0] * 2 * upsample_rate for i in
+                                      range(0, len(pred_semantic_list))]
+                    audio_frag_end_idx = [sum(audio_frag_idx[:i + 1]) for i in range(0, len(audio_frag_idx))]
+                    all_pred_semantic = torch.cat(pred_semantic_list).unsqueeze(0).unsqueeze(0).to(self.configs.device)
+                    _batch_phones = torch.cat(batch_phones).unsqueeze(0).to(self.configs.device)
+                    _batch_audio_fragment = (self.vits_model.decode(
+                        all_pred_semantic, _batch_phones, refer_audio_spec, speed=speed_factor
+                    ).detach()[0, 0, :])
+                    audio_frag_end_idx.insert(0, 0)
+                    batch_audio_fragment = [_batch_audio_fragment[audio_frag_end_idx[i - 1]:audio_frag_end_idx[i]] for i
+                                            in range(1, len(audio_frag_end_idx))]
+                else:
+                    # ## vits串行推理
+                    for i, idx in enumerate(idx_list):
+                        phones = batch_phones[i].unsqueeze(0).to(self.configs.device)
+                        _pred_semantic = (
+                            pred_semantic_list[i][-idx:].unsqueeze(0).unsqueeze(0))  # .unsqueeze(0)#mq要多unsqueeze一次
+                        audio_fragment = (self.vits_model.decode(
+                            _pred_semantic, phones, refer_audio_spec, speed=speed_factor
+                        ).detach()[0, 0, :])
+                        batch_audio_fragment.append(
+                            audio_fragment
+                        )  ###试试重建不带上prompt部分
+
+                t5 = ttime()
+                t_45 += t5 - t4
+                if return_fragment:
+                    print("%.3f\t%.3f\t%.3f\t%.3f" % (t1 - t0, t2 - t1, t4 - t3, t5 - t4))
+                    yield self.audio_postprocess([batch_audio_fragment],
+                                                 self.configs.sampling_rate,
+                                                 None,
+                                                 speed_factor,
+                                                 False,
+                                                 fragment_interval
+                                                 )
+                else:
+                    audio.append(batch_audio_fragment)
+
+                if self.stop_flag:
+                    yield self.configs.sampling_rate, np.zeros(int(self.configs.sampling_rate),
+                                                               dtype=np.int16)
+                    return
+
+            if not return_fragment:
+                # print("%.3f\t%.3f\t%.3f\t%.3f" % (t1 - t0, t2 - t1, t_34, t_45))
+                if len(audio) == 0:
+                    yield self.configs.sampling_rate, np.zeros(int(self.configs.sampling_rate),
+                                                               dtype=np.int16)
+                    return
+                yield self.audio_postprocess(audio,
+                                             self.configs.sampling_rate,
+                                             batch_index_list,
+                                             speed_factor,
+                                             split_bucket,
+                                             fragment_interval
+                                             )
+
+        except Exception as e:
+            traceback.print_exc()
+            # 必须返回一个空音频, 否则会导致显存不释放。
+            yield self.configs.sampling_rate, np.zeros(int(self.configs.sampling_rate),
+                                                       dtype=np.int16)
+            # 重置模型, 否则会导致显存释放不完全。
+            del self.t2s_model
+            del self.vits_model
+            self.t2s_model = None
+            self.vits_model = None
+            self.init_t2s_weights(self.configs.t2s_weights_path)
+            self.init_vits_weights(self.configs.vits_weights_path)
+            raise e
+        finally:
+            self.empty_cache()
+
+    def empty_cache(self):
+        try:
+            gc.collect()  # 触发gc的垃圾回收。避免内存一直增长。
+            if "cuda" in str(self.configs.device):
+                torch.cuda.empty_cache()
+            elif str(self.configs.device) == "mps":
+                torch.mps.empty_cache()
+        except:
+            pass
+
+    def audio_postprocess(self,
+                          audio: List[torch.Tensor],
+                          sr: int,
+                          batch_index_list: list = None,
+                          speed_factor: float = 1.0,
+                          split_bucket: bool = True,
+                          fragment_interval: float = 0.3
+                          ) -> Tuple[int, np.ndarray]:
+        zero_wav = torch.zeros(
+            int(self.configs.sampling_rate * fragment_interval),
+            dtype=self.precision,
+            device=self.configs.device
+        )
+
+        for i, batch in enumerate(audio):
+            for j, audio_fragment in enumerate(batch):
+                max_audio = torch.abs(audio_fragment).max()  # 简单防止16bit爆音
+                if max_audio > 1: audio_fragment /= max_audio
+                audio_fragment: torch.Tensor = torch.cat([audio_fragment, zero_wav], dim=0)
+                audio[i][j] = audio_fragment.cpu().numpy()
+
+        if split_bucket:
+            audio = self.recovery_order(audio, batch_index_list)
+        else:
+            # audio = [item for batch in audio for item in batch]
+            audio = sum(audio, [])
+
+        audio = np.concatenate(audio, 0)
+        audio = (audio * 32768).astype(np.int16)
+
+        # try:
+        #     if speed_factor != 1.0:
+        #         audio = speed_change(audio, speed=speed_factor, sr=int(sr))
+        # except Exception as e:
+        #     print(f"Failed to change speed of audio: \n{e}")
+
+        return sr, audio
+
+
+def speed_change(input_audio: np.ndarray, speed: float, sr: int):
+    # 将 NumPy 数组转换为原始 PCM 流
+    raw_audio = input_audio.astype(np.int16).tobytes()
+
+    # 设置 ffmpeg 输入流
+    input_stream = ffmpeg.input('pipe:', format='s16le', acodec='pcm_s16le', ar=str(sr), ac=1)
+
+    # 变速处理
+    output_stream = input_stream.filter('atempo', speed)
+
+    # 输出流到管道
+    out, _ = (
+        output_stream.output('pipe:', format='s16le', acodec='pcm_s16le')
+        .run(input=raw_audio, capture_stdout=True, capture_stderr=True)
+    )
+
+    # 将管道输出解码为 NumPy 数组
+    processed_audio = np.frombuffer(out, np.int16)
+
+    return processed_audio

TTS_infer_pack/TTS_mid.py

@@ -0,0 +1,1043 @@
+import gc
+import math
+import os
+import pathlib
+import random
+import sys
+import traceback
+from copy import deepcopy
+from time import time as ttime
+from typing import List, Tuple, Union
+
+import ffmpeg
+import librosa
+import numpy as np
+import torch
+import yaml
+from tqdm import tqdm
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+
+from moyoyo_tts.AR.models.t2s_lightning_module import Text2SemanticLightningModule
+from moyoyo_tts.TTS_infer_pack.TextPreprocessor import TextPreprocessor
+from moyoyo_tts.TTS_infer_pack.text_segmentation_method import splits
+from moyoyo_tts.feature_extractor.cnhubert import CNHubert
+from moyoyo_tts.module.mel_processing import spectrogram_torch
+from moyoyo_tts.module.models import SynthesizerTrn
+from moyoyo_tts.tools.i18n.i18n import I18nAuto, scan_language_list
+from moyoyo_tts.tools.my_utils import load_audio
+
+language = os.environ.get("language", "Auto")
+language = sys.argv[-1] if sys.argv[-1] in scan_language_list() else language
+i18n = I18nAuto(language=language)
+
+# configs/tts_infer.yaml
+"""
+custom:
+  bert_base_path: moyoyo_tts/pretrained_models/chinese-roberta-wwm-ext-large
+  cnhuhbert_base_path: moyoyo_tts/pretrained_models/chinese-hubert-base
+  device: cpu
+  is_half: false
+  t2s_weights_path: moyoyo_tts/pretrained_models/gsv-v2final-pretrained/s1bert25hz-5kh-longer-epoch=12-step=369668.ckpt
+  vits_weights_path: moyoyo_tts/pretrained_models/gsv-v2final-pretrained/s2G2333k.pth
+  version: v2
+default:
+  bert_base_path: moyoyo_tts/pretrained_models/chinese-roberta-wwm-ext-large
+  cnhuhbert_base_path: moyoyo_tts/pretrained_models/chinese-hubert-base
+  device: cpu
+  is_half: false
+  t2s_weights_path: moyoyo_tts/pretrained_models/s1bert25hz-2kh-longer-epoch=68e-step=50232.ckpt
+  vits_weights_path: moyoyo_tts/pretrained_models/s2G488k.pth
+  version: v1
+default_v2:
+  bert_base_path: moyoyo_tts/pretrained_models/chinese-roberta-wwm-ext-large
+  cnhuhbert_base_path: moyoyo_tts/pretrained_models/chinese-hubert-base
+  device: cpu
+  is_half: false
+  t2s_weights_path: moyoyo_tts/pretrained_models/gsv-v2final-pretrained/s1bert25hz-5kh-longer-epoch=12-step=369668.ckpt
+  vits_weights_path: moyoyo_tts/pretrained_models/gsv-v2final-pretrained/s2G2333k.pth
+  version: v2
+"""
+
+
+def set_seed(seed: int):
+    seed = int(seed)
+    seed = seed if seed != -1 else random.randrange(1 << 32)
+    #print(f"Set seed to {seed}")
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    try:
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed(seed)
+            torch.cuda.manual_seed_all(seed)
+            # torch.backends.cudnn.deterministic = True
+            # torch.backends.cudnn.benchmark = False
+            # torch.backends.cudnn.enabled = True
+            # 开启后会影响精度
+            torch.backends.cuda.matmul.allow_tf32 = False
+            torch.backends.cudnn.allow_tf32 = False
+    except:
+        pass
+    return seed
+
+
+class TTS_Config:
+    default_configs = {
+        "default": {
+            "device": "cpu",
+            "is_half": False,
+            "version": "v1",
+            "t2s_weights_path": "moyoyo_tts/pretrained_models/s1bert25hz-2kh-longer-epoch=68e-step=50232.ckpt",
+            "vits_weights_path": "moyoyo_tts/pretrained_models/s2G488k.pth",
+            "cnhuhbert_base_path": "moyoyo_tts/pretrained_models/chinese-hubert-base",
+            "bert_base_path": "moyoyo_tts/pretrained_models/chinese-roberta-wwm-ext-large",
+        },
+        "default_v2": {
+            "device": "cpu",
+            "is_half": False,
+            "version": "v2",
+            "t2s_weights_path": "moyoyo_tts/pretrained_models/gsv-v2final-pretrained/s1bert25hz-5kh-longer-epoch=12-step=369668.ckpt",
+            "vits_weights_path": "moyoyo_tts/pretrained_models/gsv-v2final-pretrained/s2G2333k.pth",
+            "cnhuhbert_base_path": "moyoyo_tts/pretrained_models/chinese-hubert-base",
+            "bert_base_path": "moyoyo_tts/pretrained_models/chinese-roberta-wwm-ext-large",
+        },
+    }
+    configs: dict = None
+    v1_languages: list = ["auto", "en", "zh", "ja", "all_zh", "all_ja"]
+    v2_languages: list = ["auto", "auto_yue", "en", "zh", "ja", "yue", "ko", "all_zh", "all_ja", "all_yue", "all_ko"]
+    languages: list = v2_languages
+
+    # "all_zh",#全部按中文识别
+    # "en",#全部按英文识别#######不变
+    # "all_ja",#全部按日文识别
+    # "all_yue",#全部按中文识别
+    # "all_ko",#全部按韩文识别
+    # "zh",#按中英混合识别####不变
+    # "ja",#按日英混合识别####不变
+    # "yue",#按粤英混合识别####不变
+    # "ko",#按韩英混合识别####不变
+    # "auto",#多语种启动切分识别语种
+    # "auto_yue",#多语种启动切分识别语种
+
+    def __init__(self, configs: Union[dict, str] = None):
+
+        if configs in ["", None]:
+            # 设置默认配置文件路径
+            configs_base_path: str = "moyoyo_tts/configs/"
+            os.makedirs(configs_base_path, exist_ok=True)
+            self.configs_path: str = os.path.join(configs_base_path, "tts_infer.yaml")
+            if not os.path.exists(self.configs_path):
+                self.save_configs()
+                print(f"Create default config file at {self.configs_path}")
+            configs: dict = deepcopy(self.default_configs)
+
+        if isinstance(configs, str):
+            self.configs_path = configs
+            configs: dict = self._load_configs(self.configs_path)
+
+        assert isinstance(configs, dict)
+        version = configs.get("version", "v2").lower()
+        assert version in ["v1", "v2"]
+        self.default_configs["default"] = configs.get("default", self.default_configs["default"])
+        self.default_configs["default_v2"] = configs.get("default_v2", self.default_configs["default_v2"])
+
+        default_config_key = "default" if version == "v1" else "default_v2"
+        self.configs: dict = configs.get("custom", deepcopy(self.default_configs[default_config_key]))
+
+        self.device = self.configs.get("device", torch.device("cpu"))
+        self.is_half = self.configs.get("is_half", False)
+        self.version = version
+        self.t2s_weights_path = self.configs.get("t2s_weights_path", None)
+        self.vits_weights_path = self.configs.get("vits_weights_path", None)
+        self.bert_base_path = self.configs.get("bert_base_path", None)
+        os.environ['bert_path'] = f'{self.bert_base_path}'
+        self.cnhuhbert_base_path = self.configs.get("cnhuhbert_base_path", None)
+        self.languages = self.v2_languages if self.version == "v2" else self.v1_languages
+
+        if (self.t2s_weights_path in [None, ""]) or (not os.path.exists(self.t2s_weights_path)):
+            self.t2s_weights_path = self.default_configs[default_config_key]['t2s_weights_path']
+            print(f"fall back to default t2s_weights_path: {self.t2s_weights_path}")
+
+        if (self.vits_weights_path in [None, ""]) or (not os.path.exists(self.vits_weights_path)):
+            self.vits_weights_path = self.default_configs[default_config_key]['vits_weights_path']
+            print(f"fall back to default vits_weights_path: {self.vits_weights_path}")
+
+        if (self.bert_base_path in [None, ""]) or (not os.path.exists(self.bert_base_path)):
+            self.bert_base_path = self.default_configs[default_config_key]['bert_base_path']
+            print(f"fall back to default bert_base_path: {self.bert_base_path}")
+
+        if (self.cnhuhbert_base_path in [None, ""]) or (not os.path.exists(self.cnhuhbert_base_path)):
+            self.cnhuhbert_base_path = self.default_configs[default_config_key]['cnhuhbert_base_path']
+            print(f"fall back to default cnhuhbert_base_path: {self.cnhuhbert_base_path}")
+
+        self.update_configs()
+
+        self.max_sec = None
+        self.hz: int = 50
+        self.semantic_frame_rate: str = "25hz"
+        self.segment_size: int = 20480
+        self.filter_length: int = 2048
+        self.sampling_rate: int = 32000
+        self.hop_length: int = 640
+        self.win_length: int = 2048
+        self.n_speakers: int = 300
+
+    def _load_configs(self, configs_path: str) -> dict:
+        if os.path.exists(configs_path):
+            ...
+        else:
+            print(i18n("路径不存在,使用默认配置"))
+            self.save_configs(configs_path)
+        with open(configs_path, 'r') as f:
+            configs = yaml.load(f, Loader=yaml.FullLoader)
+
+        return configs
+
+    def save_configs(self, configs_path: str = None) -> None:
+        configs = deepcopy(self.default_configs)
+        if self.configs is not None:
+            configs["custom"] = self.update_configs()
+
+        if configs_path is None and hasattr(self, 'configs_path'):
+            configs_path = self.configs_path
+
+        if configs_path is None:
+            return
+
+        with open(configs_path, 'w') as f:
+            yaml.dump(configs, f)
+
+    def update_configs(self):
+        self.config = {
+            "device": str(self.device),
+            "is_half": self.is_half,
+            "version": self.version,
+            "t2s_weights_path": self.t2s_weights_path,
+            "vits_weights_path": self.vits_weights_path,
+            "bert_base_path": self.bert_base_path,
+            "cnhuhbert_base_path": self.cnhuhbert_base_path,
+        }
+        return self.config
+
+    def update_version(self, version: str) -> None:
+        self.version = version
+        self.languages = self.v2_languages if self.version == "v2" else self.v1_languages
+
+    def __str__(self):
+        self.configs = self.update_configs()
+        string = "TTS Config".center(100, '-') + '\n'
+        for k, v in self.configs.items():
+            string += f"{str(k).ljust(20)}: {str(v)}\n"
+        string += "-" * 100 + '\n'
+        return string
+
+    def __repr__(self):
+        return self.__str__()
+
+    def __hash__(self):
+        return hash(self.configs_path)
+
+    def __eq__(self, other):
+        return isinstance(other, TTS_Config) and self.configs_path == other.configs_path
+
+
+class TTS:
+    def __init__(self, configs: Union[dict, str, TTS_Config]):
+        if isinstance(configs, TTS_Config):
+            self.configs = configs
+        else:
+            self.configs: TTS_Config = TTS_Config(configs)
+
+        self.t2s_model: Text2SemanticLightningModule = None
+        self.vits_model: SynthesizerTrn = None
+        self.bert_tokenizer: AutoTokenizer = None
+        self.bert_model: AutoModelForMaskedLM = None
+        self.cnhuhbert_model: CNHubert = None
+
+        self._init_models()
+
+        self.text_preprocessor: TextPreprocessor = \
+            TextPreprocessor(self.bert_model,
+                             self.bert_tokenizer,
+                             self.configs.device)
+
+        self.prompt_cache: dict = {
+            "ref_audio_path": None,
+            "prompt_semantic": None,
+            "refer_spec": [],
+            "prompt_text": None,
+            "prompt_lang": None,
+            "phones": None,
+            "bert_features": None,
+            "norm_text": None,
+            "aux_ref_audio_paths": [],
+        }
+
+        self.stop_flag: bool = False
+        self.precision: torch.dtype = torch.float16 if self.configs.is_half else torch.float32
+
+    def _init_models(self, ):
+        self.init_t2s_weights(self.configs.t2s_weights_path)
+        self.init_vits_weights(self.configs.vits_weights_path)
+        self.init_bert_weights(self.configs.bert_base_path)
+        self.init_cnhuhbert_weights(self.configs.cnhuhbert_base_path)
+        # self.enable_half_precision(self.configs.is_half)
+
+    def init_cnhuhbert_weights(self, base_path: str):
+        print(f"Loading CNHuBERT weights from {base_path}")
+        self.cnhuhbert_model = CNHubert(base_path)
+        self.cnhuhbert_model = self.cnhuhbert_model.eval()
+        self.cnhuhbert_model = self.cnhuhbert_model.to(self.configs.device)
+        if self.configs.is_half and str(self.configs.device) != "cpu":
+            self.cnhuhbert_model = self.cnhuhbert_model.half()
+
+    def init_bert_weights(self, base_path: str):
+        print(f"Loading BERT weights from {base_path}")
+        self.bert_tokenizer = AutoTokenizer.from_pretrained(base_path)
+        self.bert_model = AutoModelForMaskedLM.from_pretrained(base_path)
+        self.bert_model = self.bert_model.eval()
+        self.bert_model = self.bert_model.to(self.configs.device)
+        if self.configs.is_half and str(self.configs.device) != "cpu":
+            self.bert_model = self.bert_model.half()
+
+    def init_vits_weights(self, weights_path: str):
+        print(f"Loading VITS weights from {weights_path}")
+        self.configs.vits_weights_path = weights_path
+        dict_s2 = torch.load(weights_path, map_location=self.configs.device)
+        hps = dict_s2["config"]
+        if dict_s2['weight']['enc_p.text_embedding.weight'].shape[0] == 322:
+            self.configs.update_version("v1")
+        else:
+            self.configs.update_version("v2")
+        self.configs.save_configs()
+
+        hps["model"]["version"] = self.configs.version
+        self.configs.filter_length = hps["data"]["filter_length"]
+        self.configs.segment_size = hps["train"]["segment_size"]
+        self.configs.sampling_rate = hps["data"]["sampling_rate"]
+        self.configs.hop_length = hps["data"]["hop_length"]
+        self.configs.win_length = hps["data"]["win_length"]
+        self.configs.n_speakers = hps["data"]["n_speakers"]
+        self.configs.semantic_frame_rate = "25hz"
+        kwargs = hps["model"]
+        vits_model = SynthesizerTrn(
+            self.configs.filter_length // 2 + 1,
+            self.configs.segment_size // self.configs.hop_length,
+            n_speakers=self.configs.n_speakers,
+            **kwargs
+        )
+
+        if hasattr(vits_model, "enc_q"):
+            del vits_model.enc_q
+
+        vits_model = vits_model.to(self.configs.device)
+        vits_model = vits_model.eval()
+        vits_model.load_state_dict(dict_s2["weight"], strict=False)
+        self.vits_model = vits_model
+        if self.configs.is_half and str(self.configs.device) != "cpu":
+            self.vits_model = self.vits_model.half()
+
+    def init_t2s_weights(self, weights_path: str):
+        print(f"Loading Text2Semantic weights from {weights_path}")
+        self.configs.t2s_weights_path = weights_path
+        self.configs.save_configs()
+        self.configs.hz = 50
+        dict_s1 = torch.load(weights_path, map_location=self.configs.device)
+        config = dict_s1["config"]
+        self.configs.max_sec = config["data"]["max_sec"]
+        t2s_model = Text2SemanticLightningModule(config, "****", is_train=False)
+        t2s_model.load_state_dict(dict_s1["weight"])
+        t2s_model = t2s_model.to(self.configs.device)
+        t2s_model = t2s_model.eval()
+        self.t2s_model = t2s_model
+        if self.configs.is_half and str(self.configs.device) != "cpu":
+            self.t2s_model = self.t2s_model.half()
+
+    def enable_half_precision(self, enable: bool = True, save: bool = True):
+        '''
+            To enable half precision for the TTS model.
+            Args:
+                enable: bool, whether to enable half precision.
+
+        '''
+        if str(self.configs.device) == "cpu" and enable:
+            print("Half precision is not supported on CPU.")
+            return
+
+        self.configs.is_half = enable
+        self.precision = torch.float16 if enable else torch.float32
+        if save:
+            self.configs.save_configs()
+        if enable:
+            if self.t2s_model is not None:
+                self.t2s_model = self.t2s_model.half()
+            if self.vits_model is not None:
+                self.vits_model = self.vits_model.half()
+            if self.bert_model is not None:
+                self.bert_model = self.bert_model.half()
+            if self.cnhuhbert_model is not None:
+                self.cnhuhbert_model = self.cnhuhbert_model.half()
+        else:
+            if self.t2s_model is not None:
+                self.t2s_model = self.t2s_model.float()
+            if self.vits_model is not None:
+                self.vits_model = self.vits_model.float()
+            if self.bert_model is not None:
+                self.bert_model = self.bert_model.float()
+            if self.cnhuhbert_model is not None:
+                self.cnhuhbert_model = self.cnhuhbert_model.float()
+
+    def set_device(self, device: torch.device, save: bool = True):
+        '''
+            To set the device for all models.
+            Args:
+                device: torch.device, the device to use for all models.
+        '''
+        self.configs.device = device
+        if save:
+            self.configs.save_configs()
+        if self.t2s_model is not None:
+            self.t2s_model = self.t2s_model.to(device)
+        if self.vits_model is not None:
+            self.vits_model = self.vits_model.to(device)
+        if self.bert_model is not None:
+            self.bert_model = self.bert_model.to(device)
+        if self.cnhuhbert_model is not None:
+            self.cnhuhbert_model = self.cnhuhbert_model.to(device)
+
+    def set_ref_audio(self, ref_audio_path: str):
+        '''
+            To set the reference audio for the TTS model,
+                including the prompt_semantic and refer_spepc.
+            Args:
+                ref_audio_path: str, the path of the reference audio.
+        '''
+        self._set_prompt_semantic(ref_audio_path)
+        self._set_ref_spec(ref_audio_path)
+        self._set_ref_audio_path(ref_audio_path)
+
+    def _set_ref_audio_path(self, ref_audio_path):
+        self.prompt_cache["ref_audio_path"] = ref_audio_path
+
+    def _set_ref_spec(self, ref_audio_path):
+        spec = self._get_ref_spec(ref_audio_path)
+        if self.prompt_cache["refer_spec"] in [[], None]:
+            self.prompt_cache["refer_spec"] = [spec]
+        else:
+            self.prompt_cache["refer_spec"][0] = spec
+
+    def _get_ref_spec(self, ref_audio_path):
+        audio = load_audio(ref_audio_path, int(self.configs.sampling_rate))
+        audio = torch.FloatTensor(audio)
+        maxx = audio.abs().max()
+        if (maxx > 1): audio /= min(2, maxx)
+        audio_norm = audio
+        audio_norm = audio_norm.unsqueeze(0)
+        spec = spectrogram_torch(
+            audio_norm,
+            self.configs.filter_length,
+            self.configs.sampling_rate,
+            self.configs.hop_length,
+            self.configs.win_length,
+            center=False,
+        )
+        spec = spec.to(self.configs.device)
+        if self.configs.is_half:
+            spec = spec.half()
+        ref_audio = pathlib.Path(ref_audio_path)
+        export_dir = ref_audio.parent
+        export_path = f"{export_dir / f'refer_spec.pt'}"
+        print(f'ref spec file path: {export_path}')
+        torch.save(spec, export_path)
+        return spec
+
+    def _set_prompt_semantic(self, ref_wav_path: str):
+        zero_wav = np.zeros(
+            int(self.configs.sampling_rate * 0.3),
+            dtype=np.float16 if self.configs.is_half else np.float32,
+        )
+        with torch.no_grad():
+            wav16k, sr = librosa.load(ref_wav_path, sr=16000)
+            if (wav16k.shape[0] > 160000 or wav16k.shape[0] < 48000):
+                raise OSError(i18n("参考音频在3~10秒范围外，请更换！"))
+            wav16k = torch.from_numpy(wav16k)
+            zero_wav_torch = torch.from_numpy(zero_wav)
+            wav16k = wav16k.to(self.configs.device)
+            zero_wav_torch = zero_wav_torch.to(self.configs.device)
+            if self.configs.is_half:
+                wav16k = wav16k.half()
+                zero_wav_torch = zero_wav_torch.half()
+
+            wav16k = torch.cat([wav16k, zero_wav_torch])
+            hubert_feature = self.cnhuhbert_model.model(wav16k.unsqueeze(0))[
+                "last_hidden_state"
+            ].transpose(
+                1, 2
+            )  # .float()
+            codes = self.vits_model.extract_latent(hubert_feature)
+
+            prompt_semantic = codes[0, 0].to(self.configs.device)
+            self.prompt_cache["prompt_semantic"] = prompt_semantic
+
+            ref_audio = pathlib.Path(ref_wav_path)
+            export_dir = ref_audio.parent
+            export_path = f"{export_dir / f'prompt_semantic.pt'}"
+            print(f'prompt semantic file path: {export_path}')
+            torch.save(prompt_semantic, export_path)
+
+    def batch_sequences(self, sequences: List[torch.Tensor], axis: int = 0, pad_value: int = 0, max_length: int = None):
+        seq = sequences[0]
+        ndim = seq.dim()
+        if axis < 0:
+            axis += ndim
+        dtype: torch.dtype = seq.dtype
+        pad_value = torch.tensor(pad_value, dtype=dtype)
+        seq_lengths = [seq.shape[axis] for seq in sequences]
+        if max_length is None:
+            max_length = max(seq_lengths)
+        else:
+            max_length = max(seq_lengths) if max_length < max(seq_lengths) else max_length
+
+        padded_sequences = []
+        for seq, length in zip(sequences, seq_lengths):
+            padding = [0] * axis + [0, max_length - length] + [0] * (ndim - axis - 1)
+            padded_seq = torch.nn.functional.pad(seq, padding, value=pad_value)
+            padded_sequences.append(padded_seq)
+        batch = torch.stack(padded_sequences)
+        return batch
+
+    def to_batch(self, data: list,
+                 prompt_data: dict = None,
+                 batch_size: int = 5,
+                 threshold: float = 0.75,
+                 split_bucket: bool = True,
+                 device: torch.device = torch.device("cpu"),
+                 precision: torch.dtype = torch.float32,
+                 ):
+        _data: list = []
+        index_and_len_list = []
+        for idx, item in enumerate(data):
+            norm_text_len = len(item["norm_text"])
+            index_and_len_list.append([idx, norm_text_len])
+
+        batch_index_list = []
+        if split_bucket:
+            index_and_len_list.sort(key=lambda x: x[1])
+            index_and_len_list = np.array(index_and_len_list, dtype=np.int64)
+
+            batch_index_list_len = 0
+            pos = 0
+            while pos < index_and_len_list.shape[0]:
+                # batch_index_list.append(index_and_len_list[pos:min(pos+batch_size,len(index_and_len_list))])
+                pos_end = min(pos + batch_size, index_and_len_list.shape[0])
+                while pos < pos_end:
+                    batch = index_and_len_list[pos:pos_end, 1].astype(np.float32)
+                    score = batch[(pos_end - pos) // 2] / (batch.mean() + 1e-8)
+                    if (score >= threshold) or (pos_end - pos == 1):
+                        batch_index = index_and_len_list[pos:pos_end, 0].tolist()
+                        batch_index_list_len += len(batch_index)
+                        batch_index_list.append(batch_index)
+                        pos = pos_end
+                        break
+                    pos_end = pos_end - 1
+
+            assert batch_index_list_len == len(data)
+
+        else:
+            for i in range(len(data)):
+                if i % batch_size == 0:
+                    batch_index_list.append([])
+                batch_index_list[-1].append(i)
+
+        for batch_idx, index_list in enumerate(batch_index_list):
+            item_list = [data[idx] for idx in index_list]
+            phones_list = []
+            phones_len_list = []
+            # bert_features_list = []
+            all_phones_list = []
+            all_phones_len_list = []
+            all_bert_features_list = []
+            norm_text_batch = []
+            all_bert_max_len = 0
+            all_phones_max_len = 0
+            for item in item_list:
+                if prompt_data is not None:
+                    all_bert_features = torch.cat([prompt_data["bert_features"], item["bert_features"]], 1) \
+                        .to(dtype=precision, device=device)
+                    all_phones = torch.LongTensor(prompt_data["phones"] + item["phones"]).to(device)
+                    phones = torch.LongTensor(item["phones"]).to(device)
+                    # norm_text = prompt_data["norm_text"]+item["norm_text"]
+                else:
+                    all_bert_features = item["bert_features"] \
+                        .to(dtype=precision, device=device)
+                    phones = torch.LongTensor(item["phones"]).to(device)
+                    all_phones = phones
+                    # norm_text = item["norm_text"]
+
+                all_bert_max_len = max(all_bert_max_len, all_bert_features.shape[-1])
+                all_phones_max_len = max(all_phones_max_len, all_phones.shape[-1])
+
+                phones_list.append(phones)
+                phones_len_list.append(phones.shape[-1])
+                all_phones_list.append(all_phones)
+                all_phones_len_list.append(all_phones.shape[-1])
+                all_bert_features_list.append(all_bert_features)
+                norm_text_batch.append(item["norm_text"])
+
+            phones_batch = phones_list
+            all_phones_batch = all_phones_list
+            all_bert_features_batch = all_bert_features_list
+
+            max_len = max(all_bert_max_len, all_phones_max_len)
+            # phones_batch = self.batch_sequences(phones_list, axis=0, pad_value=0, max_length=max_len)
+            #### 直接对phones和bert_features进行pad。（padding策略会影响T2S模型生成的结果，但不直接影响复读概率。影响复读概率的主要因素是mask的策略）
+            # all_phones_batch = self.batch_sequences(all_phones_list, axis=0, pad_value=0, max_length=max_len)
+            # all_bert_features_batch = all_bert_features_list
+            # all_bert_features_batch = torch.zeros((len(all_bert_features_list), 1024, max_len), dtype=precision, device=device)
+            # for idx, item in enumerate(all_bert_features_list):
+            #     all_bert_features_batch[idx, :, : item.shape[-1]] = item
+
+            # #### 先对phones进行embedding、对bert_features进行project，再pad到相同长度，（padding策略会影响T2S模型生成的结果，但不直接影响复读概率。影响复读概率的主要因素是mask的策略）
+            # all_phones_list = [self.t2s_model.model.ar_text_embedding(item.to(self.t2s_model.device)) for item in all_phones_list]
+            # all_phones_list = [F.pad(item,(0,0,0,max_len-item.shape[0]),value=0) for item in all_phones_list]
+            # all_phones_batch = torch.stack(all_phones_list, dim=0)
+
+            # all_bert_features_list = [self.t2s_model.model.bert_proj(item.to(self.t2s_model.device).transpose(0, 1)) for item in all_bert_features_list]
+            # all_bert_features_list = [F.pad(item,(0,0,0,max_len-item.shape[0]), value=0) for item in all_bert_features_list]
+            # all_bert_features_batch = torch.stack(all_bert_features_list, dim=0)
+
+            batch = {
+                "phones": phones_batch,
+                "phones_len": torch.LongTensor(phones_len_list).to(device),
+                "all_phones": all_phones_batch,
+                "all_phones_len": torch.LongTensor(all_phones_len_list).to(device),
+                "all_bert_features": all_bert_features_batch,
+                "norm_text": norm_text_batch,
+                "max_len": max_len,
+            }
+            _data.append(batch)
+
+        return _data, batch_index_list
+
+    def recovery_order(self, data: list, batch_index_list: list) -> list:
+        '''
+        Recovery the order of the audio according to the batch_index_list.
+
+        Args:
+            data (List[list(np.ndarray)]): the out of order audio .
+            batch_index_list (List[list[int]]): the batch index list.
+
+        Returns:
+            list (List[np.ndarray]): the data in the original order.
+        '''
+        length = len(sum(batch_index_list, []))
+        _data = [None] * length
+        for i, index_list in enumerate(batch_index_list):
+            for j, index in enumerate(index_list):
+                _data[index] = data[i][j]
+        return _data
+
+    def stop(self, ):
+        '''
+        Stop the inference process.
+        '''
+        self.stop_flag = True
+
+    @torch.no_grad()
+    def run(self, inputs: dict):
+        """
+        Text to speech inference.
+
+        Args:
+            inputs (dict):
+                {
+                    "text": "",                   # str.(required) text to be synthesized
+                    "text_lang: "",               # str.(required) language of the text to be synthesized
+                    "ref_audio_path": "",         # str.(required) reference audio path
+                    "aux_ref_audio_paths": [],    # list.(optional) auxiliary reference audio paths for multi-speaker tone fusion
+                    "prompt_text": "",            # str.(optional) prompt text for the reference audio
+                    "prompt_lang": "",            # str.(required) language of the prompt text for the reference audio
+                    "top_k": 5,                   # int. top k sampling
+                    "top_p": 1,                   # float. top p sampling
+                    "temperature": 1,             # float. temperature for sampling
+                    "text_split_method": "cut0",  # str. text split method, see text_segmentation_method.py for details.
+                    "batch_size": 1,              # int. batch size for inference
+                    "batch_threshold": 0.75,      # float. threshold for batch splitting.
+                    "split_bucket: True,          # bool. whether to split the batch into multiple buckets.
+                    "return_fragment": False,     # bool. step by step return the audio fragment.
+                    "speed_factor":1.0,           # float. control the speed of the synthesized audio.
+                    "fragment_interval":0.3,      # float. to control the interval of the audio fragment.
+                    "seed": -1,                   # int. random seed for reproducibility.
+                    "parallel_infer": True,       # bool. whether to use parallel inference.
+                    "repetition_penalty": 1.35    # float. repetition penalty for T2S model.
+                }
+        returns:
+            Tuple[int, np.ndarray]: sampling rate and audio data.
+        """
+        ########## variables initialization ###########
+        self.stop_flag: bool = False
+        text: str = inputs.get("text", "")
+        text_lang: str = inputs.get("text_lang", "")
+        ref_audio_path: str = inputs.get("ref_audio_path", "")
+        aux_ref_audio_paths: list = inputs.get("aux_ref_audio_paths", [])
+        prompt_text: str = inputs.get("prompt_text", "")
+        prompt_lang: str = inputs.get("prompt_lang", "")
+        top_k: int = inputs.get("top_k", 5)
+        top_p: float = inputs.get("top_p", 1)
+        temperature: float = inputs.get("temperature", 1)
+        text_split_method: str = inputs.get("text_split_method", "cut0")
+        batch_size = inputs.get("batch_size", 1)
+        batch_threshold = inputs.get("batch_threshold", 0.75)
+        speed_factor = inputs.get("speed_factor", 1.0)
+        split_bucket = inputs.get("split_bucket", True)
+        return_fragment = inputs.get("return_fragment", False)
+        fragment_interval = inputs.get("fragment_interval", 0.3)
+        seed = inputs.get("seed", -1)
+        seed = -1 if seed in ["", None] else seed
+        actual_seed = set_seed(seed)
+        parallel_infer = inputs.get("parallel_infer", True)
+        repetition_penalty = inputs.get("repetition_penalty", 1.35)
+
+        if parallel_infer:
+            #print(i18n("并行推理模式已开启"))
+            self.t2s_model.model.infer_panel = self.t2s_model.model.infer_panel_batch_infer
+        else:
+            #print(i18n("并行推理模式已关闭"))
+            self.t2s_model.model.infer_panel = self.t2s_model.model.infer_panel_naive_batched
+
+        if return_fragment:
+            #print(i18n("分段返回模式已开启"))
+            if split_bucket:
+                split_bucket = False
+                #print(i18n("分段返回模式不支持分桶处理，已自动关闭分桶处理"))
+
+        if split_bucket and speed_factor == 1.0:
+            print(i18n("分桶处理模式已开启"))
+        elif speed_factor != 1.0:
+            #print(i18n("语速调节不支持分桶处理，已自动关闭分桶处理"))
+            split_bucket = False
+        else:
+            print(i18n("分桶处理模式已关闭"))
+
+        if fragment_interval < 0.01:
+            fragment_interval = 0.01
+            #print(i18n("分段间隔过小，已自动设置为0.01"))
+
+        no_prompt_text = False
+        if prompt_text in [None, ""]:
+            no_prompt_text = True
+
+        assert text_lang in self.configs.languages
+        if not no_prompt_text:
+            assert prompt_lang in self.configs.languages
+
+        if ref_audio_path in [None, ""] and \
+                ((self.prompt_cache["prompt_semantic"] is None) or (self.prompt_cache["refer_spec"] in [None, []])):
+            raise ValueError(
+                "ref_audio_path cannot be empty, when the reference audio is not set using set_ref_audio()")
+
+        ###### setting reference audio and prompt text preprocessing ########
+        t0 = ttime()
+        if (ref_audio_path is not None) and (ref_audio_path != self.prompt_cache["ref_audio_path"]):
+            if not os.path.exists(ref_audio_path):
+                raise ValueError(f"{ref_audio_path} not exists")
+            self.set_ref_audio(ref_audio_path)
+
+        aux_ref_audio_paths = aux_ref_audio_paths if aux_ref_audio_paths is not None else []
+        paths = set(aux_ref_audio_paths) & set(self.prompt_cache["aux_ref_audio_paths"])
+        if not (len(list(paths)) == len(aux_ref_audio_paths) == len(self.prompt_cache["aux_ref_audio_paths"])):
+            self.prompt_cache["aux_ref_audio_paths"] = aux_ref_audio_paths
+            self.prompt_cache["refer_spec"] = [self.prompt_cache["refer_spec"][0]]
+            for path in aux_ref_audio_paths:
+                if path in [None, ""]:
+                    continue
+                if not os.path.exists(path):
+                    print(i18n("音频文件不存在，跳过：{}").format(path))
+                    continue
+                self.prompt_cache["refer_spec"].append(self._get_ref_spec(path))
+
+        if not no_prompt_text:
+            prompt_text = prompt_text.strip("\n")
+            if (prompt_text[-1] not in splits): prompt_text += "。" if prompt_lang != "en" else "."
+            #print(i18n("实际输入的参考文本:"), prompt_text)
+            if self.prompt_cache["prompt_text"] != prompt_text:
+                self.prompt_cache["prompt_text"] = prompt_text
+                self.prompt_cache["prompt_lang"] = prompt_lang
+                phones, bert_features, norm_text = \
+                    self.text_preprocessor.segment_and_extract_feature_for_text(
+                        prompt_text,
+                        prompt_lang,
+                        self.configs.version)
+                self.prompt_cache["phones"] = phones
+                self.prompt_cache["bert_features"] = bert_features
+                self.prompt_cache["norm_text"] = norm_text
+
+        ###### text preprocessing ########
+        t1 = ttime()
+        data: list = None
+        if not return_fragment:
+            data = self.text_preprocessor.preprocess(text, text_lang, text_split_method, self.configs.version)
+            if len(data) == 0:
+                yield self.configs.sampling_rate, np.zeros(int(self.configs.sampling_rate),
+                                                           dtype=np.int16)
+                return
+
+            batch_index_list: list = None
+            data, batch_index_list = self.to_batch(data,
+                                                   prompt_data=self.prompt_cache if not no_prompt_text else None,
+                                                   batch_size=batch_size,
+                                                   threshold=batch_threshold,
+                                                   split_bucket=split_bucket,
+                                                   device=self.configs.device,
+                                                   precision=self.precision
+                                                   )
+        else:
+            #print(i18n("############ 切分文本 ############"))
+            texts = self.text_preprocessor.pre_seg_text(text, text_lang, text_split_method)
+            data = []
+            for i in range(len(texts)):
+                if i % batch_size == 0:
+                    data.append([])
+                data[-1].append(texts[i])
+
+            def make_batch(batch_texts):
+                batch_data = []
+                #print(i18n("############ 提取文本Bert特征 ############"))
+                for text in tqdm(batch_texts):
+                    phones, bert_features, norm_text = self.text_preprocessor.segment_and_extract_feature_for_text(text,
+                                                                                                                   text_lang,
+                                                                                                                   self.configs.version)
+                    if phones is None:
+                        continue
+                    res = {
+                        "phones": phones,
+                        "bert_features": bert_features,
+                        "norm_text": norm_text,
+                    }
+                    batch_data.append(res)
+                if len(batch_data) == 0:
+                    return None
+                batch, _ = self.to_batch(batch_data,
+                                         prompt_data=self.prompt_cache if not no_prompt_text else None,
+                                         batch_size=batch_size,
+                                         threshold=batch_threshold,
+                                         split_bucket=False,
+                                         device=self.configs.device,
+                                         precision=self.precision
+                                         )
+                return batch[0]
+
+        t2 = ttime()
+        try:
+            #print("############ 推理 ############")
+            ###### inference ######
+            t_34 = 0.0
+            t_45 = 0.0
+            audio = []
+            for item in data:
+                t3 = ttime()
+                if return_fragment:
+                    item = make_batch(item)
+                    if item is None:
+                        continue
+
+                batch_phones: List[torch.LongTensor] = item["phones"]
+                # batch_phones:torch.LongTensor = item["phones"]
+                batch_phones_len: torch.LongTensor = item["phones_len"]
+                all_phoneme_ids: torch.LongTensor = item["all_phones"]
+                all_phoneme_lens: torch.LongTensor = item["all_phones_len"]
+                all_bert_features: torch.LongTensor = item["all_bert_features"]
+                norm_text: str = item["norm_text"]
+                max_len = item["max_len"]
+
+                #print(i18n("前端处理后的文本(每句):"), norm_text)
+                if no_prompt_text:
+                    prompt = None
+                else:
+                    prompt = self.prompt_cache["prompt_semantic"].expand(len(all_phoneme_ids), -1).to(
+                        self.configs.device)
+
+                pred_semantic_list, idx_list = self.t2s_model.model.infer_panel(
+                    all_phoneme_ids,
+                    all_phoneme_lens,
+                    prompt,
+                    all_bert_features,
+                    # prompt_phone_len=ph_offset,
+                    top_k=top_k,
+                    top_p=top_p,
+                    temperature=temperature,
+                    early_stop_num=self.configs.hz * self.configs.max_sec,
+                    max_len=max_len,
+                    repetition_penalty=repetition_penalty,
+                )
+                t4 = ttime()
+                t_34 += t4 - t3
+
+                refer_audio_spec: torch.Tensor = [item.to(dtype=self.precision, device=self.configs.device) for item in
+                                                  self.prompt_cache["refer_spec"]]
+
+                batch_audio_fragment = []
+
+                # ## vits并行推理 method 1
+                # pred_semantic_list = [item[-idx:] for item, idx in zip(pred_semantic_list, idx_list)]
+                # pred_semantic_len = torch.LongTensor([item.shape[0] for item in pred_semantic_list]).to(self.configs.device)
+                # pred_semantic = self.batch_sequences(pred_semantic_list, axis=0, pad_value=0).unsqueeze(0)
+                # max_len = 0
+                # for i in range(0, len(batch_phones)):
+                #     max_len = max(max_len, batch_phones[i].shape[-1])
+                # batch_phones = self.batch_sequences(batch_phones, axis=0, pad_value=0, max_length=max_len)
+                # batch_phones = batch_phones.to(self.configs.device)
+                # batch_audio_fragment = (self.vits_model.batched_decode(
+                #         pred_semantic, pred_semantic_len, batch_phones, batch_phones_len,refer_audio_spec
+                #     ))
+
+                if speed_factor == 1.0:
+                    # ## vits并行推理 method 2
+                    pred_semantic_list = [item[-idx:] for item, idx in zip(pred_semantic_list, idx_list)]
+                    upsample_rate = math.prod(self.vits_model.upsample_rates)
+                    audio_frag_idx = [pred_semantic_list[i].shape[0] * 2 * upsample_rate for i in
+                                      range(0, len(pred_semantic_list))]
+                    audio_frag_end_idx = [sum(audio_frag_idx[:i + 1]) for i in range(0, len(audio_frag_idx))]
+                    all_pred_semantic = torch.cat(pred_semantic_list).unsqueeze(0).unsqueeze(0).to(self.configs.device)
+                    _batch_phones = torch.cat(batch_phones).unsqueeze(0).to(self.configs.device)
+                    _batch_audio_fragment = (self.vits_model.decode(
+                        all_pred_semantic, _batch_phones, refer_audio_spec, speed=speed_factor
+                    ).detach()[0, 0, :])
+                    audio_frag_end_idx.insert(0, 0)
+                    batch_audio_fragment = [_batch_audio_fragment[audio_frag_end_idx[i - 1]:audio_frag_end_idx[i]] for i
+                                            in range(1, len(audio_frag_end_idx))]
+                else:
+                    # ## vits串行推理
+                    for i, idx in enumerate(idx_list):
+                        phones = batch_phones[i].unsqueeze(0).to(self.configs.device)
+                        _pred_semantic = (
+                            pred_semantic_list[i][-idx:].unsqueeze(0).unsqueeze(0))  # .unsqueeze(0)#mq要多unsqueeze一次
+                        audio_fragment = (self.vits_model.decode(
+                            _pred_semantic, phones, refer_audio_spec, speed=speed_factor
+                        ).detach()[0, 0, :])
+                        batch_audio_fragment.append(
+                            audio_fragment
+                        )  ###试试重建不带上prompt部分
+
+                t5 = ttime()
+                t_45 += t5 - t4
+                if return_fragment:
+                    print("%.3f\t%.3f\t%.3f\t%.3f" % (t1 - t0, t2 - t1, t4 - t3, t5 - t4))
+                    yield self.audio_postprocess([batch_audio_fragment],
+                                                 self.configs.sampling_rate,
+                                                 None,
+                                                 speed_factor,
+                                                 False,
+                                                 fragment_interval
+                                                 )
+                else:
+                    audio.append(batch_audio_fragment)
+
+                if self.stop_flag:
+                    yield self.configs.sampling_rate, np.zeros(int(self.configs.sampling_rate),
+                                                               dtype=np.int16)
+                    return
+
+            if not return_fragment:
+                # print("%.3f\t%.3f\t%.3f\t%.3f" % (t1 - t0, t2 - t1, t_34, t_45))
+                if len(audio) == 0:
+                    yield self.configs.sampling_rate, np.zeros(int(self.configs.sampling_rate),
+                                                               dtype=np.int16)
+                    return
+                yield self.audio_postprocess(audio,
+                                             self.configs.sampling_rate,
+                                             batch_index_list,
+                                             speed_factor,
+                                             split_bucket,
+                                             fragment_interval
+                                             )
+
+        except Exception as e:
+            traceback.print_exc()
+            # 必须返回一个空音频, 否则会导致显存不释放。
+            yield self.configs.sampling_rate, np.zeros(int(self.configs.sampling_rate),
+                                                       dtype=np.int16)
+            # 重置模型, 否则会导致显存释放不完全。
+            del self.t2s_model
+            del self.vits_model
+            self.t2s_model = None
+            self.vits_model = None
+            self.init_t2s_weights(self.configs.t2s_weights_path)
+            self.init_vits_weights(self.configs.vits_weights_path)
+            raise e
+        finally:
+            self.empty_cache()
+
+    def empty_cache(self):
+        try:
+            gc.collect()  # 触发gc的垃圾回收。避免内存一直增长。
+            if "cuda" in str(self.configs.device):
+                torch.cuda.empty_cache()
+            elif str(self.configs.device) == "mps":
+                torch.mps.empty_cache()
+        except:
+            pass
+
+    def audio_postprocess(self,
+                          audio: List[torch.Tensor],
+                          sr: int,
+                          batch_index_list: list = None,
+                          speed_factor: float = 1.0,
+                          split_bucket: bool = True,
+                          fragment_interval: float = 0.3
+                          ) -> Tuple[int, np.ndarray]:
+        zero_wav = torch.zeros(
+            int(self.configs.sampling_rate * fragment_interval),
+            dtype=self.precision,
+            device=self.configs.device
+        )
+
+        for i, batch in enumerate(audio):
+            for j, audio_fragment in enumerate(batch):
+                max_audio = torch.abs(audio_fragment).max()  # 简单防止16bit爆音
+                if max_audio > 1: audio_fragment /= max_audio
+                audio_fragment: torch.Tensor = torch.cat([audio_fragment, zero_wav], dim=0)
+                audio[i][j] = audio_fragment.cpu().numpy()
+
+        if split_bucket:
+            audio = self.recovery_order(audio, batch_index_list)
+        else:
+            # audio = [item for batch in audio for item in batch]
+            audio = sum(audio, [])
+
+        audio = np.concatenate(audio, 0)
+        audio = (audio * 32768).astype(np.int16)
+
+        # try:
+        #     if speed_factor != 1.0:
+        #         audio = speed_change(audio, speed=speed_factor, sr=int(sr))
+        # except Exception as e:
+        #     print(f"Failed to change speed of audio: \n{e}")
+
+        return sr, audio
+
+
+def speed_change(input_audio: np.ndarray, speed: float, sr: int):
+    # 将 NumPy 数组转换为原始 PCM 流
+    raw_audio = input_audio.astype(np.int16).tobytes()
+
+    # 设置 ffmpeg 输入流
+    input_stream = ffmpeg.input('pipe:', format='s16le', acodec='pcm_s16le', ar=str(sr), ac=1)
+
+    # 变速处理
+    output_stream = input_stream.filter('atempo', speed)
+
+    # 输出流到管道
+    out, _ = (
+        output_stream.output('pipe:', format='s16le', acodec='pcm_s16le')
+        .run(input=raw_audio, capture_stdout=True, capture_stderr=True)
+    )
+
+    # 将管道输出解码为 NumPy 数组
+    processed_audio = np.frombuffer(out, np.int16)
+
+    return processed_audio

TTS_infer_pack/TextPreprocessor.py

@@ -0,0 +1,247 @@
+
+import os
+import re
+import sys
+from typing import Dict, List, Tuple
+
+import LangSegment.LangSegment as LangSegment
+import torch
+from tqdm import tqdm
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+
+from moyoyo_tts.TTS_infer_pack.text_segmentation_method import split_big_text, splits, get_method as get_seg_method
+from moyoyo_tts.text import chinese
+from moyoyo_tts.text import cleaned_text_to_sequence
+from moyoyo_tts.text.cleaner import clean_text
+from moyoyo_tts.tools.i18n.i18n import I18nAuto, scan_language_list
+
+
+language=os.environ.get("language","Auto")
+language=sys.argv[-1] if sys.argv[-1] in scan_language_list() else language
+i18n = I18nAuto(language=language)
+punctuation = set(['!', '?', '…', ',', '.', '-'," "])
+
+def get_first(text:str) -> str:
+    pattern = "[" + "".join(re.escape(sep) for sep in splits) + "]"
+    text = re.split(pattern, text)[0].strip()
+    return text
+
+def merge_short_text_in_array(texts:str, threshold:int) -> list:
+    if (len(texts)) < 2:
+        return texts
+    result = []
+    text = ""
+    for ele in texts:
+        text += ele
+        if len(text) >= threshold:
+            result.append(text)
+            text = ""
+    if (len(text) > 0):
+        if len(result) == 0:
+            result.append(text)
+        else:
+            result[len(result) - 1] += text
+    return result
+
+
+
+
+
+class TextPreprocessor:
+    def __init__(self, bert_model:AutoModelForMaskedLM,
+                 tokenizer:AutoTokenizer, device:torch.device):
+        self.bert_model = bert_model
+        self.tokenizer = tokenizer
+        self.device = device
+
+    def preprocess(self, text:str, lang:str, text_split_method:str, version:str="v2")->List[Dict]:
+        #print(i18n("############ 切分文本 ############"))
+        text = self.replace_consecutive_punctuation(text)
+        texts = self.pre_seg_text(text, lang, text_split_method)
+        result = []
+        #print(i18n("############ 提取文本Bert特征 ############"))
+        # for text in tqdm(texts):
+        for text in texts:
+            phones, bert_features, norm_text = self.segment_and_extract_feature_for_text(text, lang, version)
+            if phones is None or norm_text=="":
+                continue
+            res={
+                "phones": phones,
+                "bert_features": bert_features,
+                "norm_text": norm_text,
+            }
+            result.append(res)
+        return result
+
+    def pre_seg_text(self, text:str, lang:str, text_split_method:str):
+        text = text.strip("\n")
+        if len(text) == 0:
+            return []
+        if (text[0] not in splits and len(get_first(text)) < 4):
+            text = "。" + text if lang != "en" else "." + text
+        #print(i18n("实际输入的目标文本:"))
+        #print(text)
+
+        seg_method = get_seg_method(text_split_method)
+        text = seg_method(text)
+
+        while "\n\n" in text:
+            text = text.replace("\n\n", "\n")
+
+        _texts = text.split("\n")
+        _texts = self.filter_text(_texts)
+        _texts = merge_short_text_in_array(_texts, 5)
+        texts = []
+
+
+        for text in _texts:
+            # 解决输入目标文本的空行导致报错的问题
+            if (len(text.strip()) == 0):
+               continue
+            if not re.sub("\W+", "", text):
+                # 检测一下，如果是纯符号，就跳过。
+                continue
+            if (text[-1] not in splits): text += "。" if lang != "en" else "."
+
+            # 解决句子过长导致Bert报错的问题
+            if (len(text) > 510):
+                texts.extend(split_big_text(text))
+            else:
+                texts.append(text)
+
+        #print(i18n("实际输入的目标文本(切句后):"))
+        #print(texts)
+        return texts
+
+    def segment_and_extract_feature_for_text(self, text:str, language:str, version:str="v1")->Tuple[list, torch.Tensor, str]:
+        return self.get_phones_and_bert(text, language, version)
+
+    def get_phones_and_bert(self, text:str, language:str, version:str, final:bool=False):
+        if language in {"en", "all_zh", "all_ja", "all_ko", "all_yue"}:
+            language = language.replace("all_","")
+            if language == "en":
+                LangSegment.setfilters(["en"])
+                formattext = " ".join(tmp["text"] for tmp in LangSegment.getTexts(text))
+            else:
+                # 因无法区别中日韩文汉字,以用户输入为准
+                formattext = text
+            while "  " in formattext:
+                formattext = formattext.replace("  ", " ")
+            if language == "zh":
+                if re.search(r'[A-Za-z]', formattext):
+                    formattext = re.sub(r'[a-z]', lambda x: x.group(0).upper(), formattext)
+                    formattext = chinese.mix_text_normalize(formattext)
+                    return self.get_phones_and_bert(formattext,"zh",version)
+                else:
+                    phones, word2ph, norm_text = self.clean_text_inf(formattext, language, version)
+                    bert = self.get_bert_feature(norm_text, word2ph).to(self.device)
+            elif language == "yue" and re.search(r'[A-Za-z]', formattext):
+                    formattext = re.sub(r'[a-z]', lambda x: x.group(0).upper(), formattext)
+                    formattext = chinese.mix_text_normalize(formattext)
+                    return self.get_phones_and_bert(formattext,"yue",version)
+            else:
+                phones, word2ph, norm_text = self.clean_text_inf(formattext, language, version)
+                bert = torch.zeros(
+                    (1024, len(phones)),
+                    dtype=torch.float32,
+                ).to(self.device)
+        elif language in {"zh", "ja", "ko", "yue", "auto", "auto_yue"}:
+            textlist=[]
+            langlist=[]
+            LangSegment.setfilters(["zh","ja","en","ko"])
+            if language == "auto":
+                for tmp in LangSegment.getTexts(text):
+                    langlist.append(tmp["lang"])
+                    textlist.append(tmp["text"])
+            elif language == "auto_yue":
+                for tmp in LangSegment.getTexts(text):
+                    if tmp["lang"] == "zh":
+                        tmp["lang"] = "yue"
+                    langlist.append(tmp["lang"])
+                    textlist.append(tmp["text"])
+            else:
+                for tmp in LangSegment.getTexts(text):
+                    if tmp["lang"] == "en":
+                        langlist.append(tmp["lang"])
+                    else:
+                        # 因无法区别中日韩文汉字,以用户输入为准
+                        langlist.append(language)
+                    textlist.append(tmp["text"])
+            # print(textlist)
+            # print(langlist)
+            phones_list = []
+            bert_list = []
+            norm_text_list = []
+            for i in range(len(textlist)):
+                lang = langlist[i]
+                phones, word2ph, norm_text = self.clean_text_inf(textlist[i], lang, version)
+                bert = self.get_bert_inf(phones, word2ph, norm_text, lang)
+                phones_list.append(phones)
+                norm_text_list.append(norm_text)
+                bert_list.append(bert)
+            bert = torch.cat(bert_list, dim=1)
+            phones = sum(phones_list, [])
+            norm_text = ''.join(norm_text_list)
+
+        if not final and len(phones) < 6:
+            return self.get_phones_and_bert("." + text,language,version,final=True)
+
+        return phones, bert, norm_text
+
+
+    def get_bert_feature(self, text:str, word2ph:list)->torch.Tensor:
+        with torch.no_grad():
+            inputs = self.tokenizer(text, return_tensors="pt")
+            for i in inputs:
+                inputs[i] = inputs[i].to(self.device)
+            res = self.bert_model(**inputs, output_hidden_states=True)
+            res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu()[1:-1]
+        assert len(word2ph) == len(text)
+        phone_level_feature = []
+        for i in range(len(word2ph)):
+            repeat_feature = res[i].repeat(word2ph[i], 1)
+            phone_level_feature.append(repeat_feature)
+        phone_level_feature = torch.cat(phone_level_feature, dim=0)
+        return phone_level_feature.T
+
+    def clean_text_inf(self, text:str, language:str, version:str="v2"):
+        phones, word2ph, norm_text = clean_text(text, language, version)
+        phones = cleaned_text_to_sequence(phones, version)
+        return phones, word2ph, norm_text
+
+    def get_bert_inf(self, phones:list, word2ph:list, norm_text:str, language:str):
+        language=language.replace("all_","")
+        if language == "zh":
+            feature = self.get_bert_feature(norm_text, word2ph).to(self.device)
+        else:
+            feature = torch.zeros(
+                (1024, len(phones)),
+                dtype=torch.float32,
+            ).to(self.device)
+
+        return feature
+
+
+    def filter_text(self,texts):
+        _text=[]
+        if all(text in [None, " ", "\n",""] for text in texts):
+            raise ValueError(i18n("请输入有效文本"))
+        for text in texts:
+            if text in  [None, " ", ""]:
+                pass
+            else:
+                _text.append(text)
+        return _text
+
+
+    def replace_consecutive_punctuation(self,text):
+        punctuations = ''.join(re.escape(p) for p in punctuation)
+        pattern = f'([{punctuations}])([{punctuations}])+'
+        result = re.sub(pattern, r'\1', text)
+        return result
+
+
+

TTS_infer_pack/__init__.py

@@ -0,0 +1 @@
+from . import TTS, text_segmentation_method

TTS_infer_pack/text_segmentation_method.py

@@ -0,0 +1,173 @@
+
+
+
+
+import re
+from typing import Callable
+
+punctuation = set(['!', '?', '…', ',', '.', '-'," "])
+METHODS = dict()
+
+def get_method(name:str)->Callable:
+    method = METHODS.get(name, None)
+    if method is None:
+        raise ValueError(f"Method {name} not found")
+    return method
+
+def get_method_names()->list:
+    return list(METHODS.keys())
+
+def register_method(name):
+    def decorator(func):
+        METHODS[name] = func
+        return func
+    return decorator
+
+splits = {"，", "。", "？", "！", ",", ".", "?", "!", "~", ":", "：", "—", "…", }
+
+def split_big_text(text, max_len=510):
+    # 定义全角和半角标点符号
+    punctuation = "".join(splits)
+
+    # 切割文本
+    segments = re.split('([' + punctuation + '])', text)
+    
+    # 初始化结果列表和当前片段
+    result = []
+    current_segment = ''
+    
+    for segment in segments:
+        # 如果当前片段加上新的片段长度超过max_len，就将当前片段加入结果列表，并重置当前片段
+        if len(current_segment + segment) > max_len:
+            result.append(current_segment)
+            current_segment = segment
+        else:
+            current_segment += segment
+    
+    # 将最后一个片段加入结果列表
+    if current_segment:
+        result.append(current_segment)
+    
+    return result
+
+
+
+def split(todo_text):
+    todo_text = todo_text.replace("……", "。").replace("——", "，")
+    if todo_text[-1] not in splits:
+        todo_text += "。"
+    i_split_head = i_split_tail = 0
+    len_text = len(todo_text)
+    todo_texts = []
+    while 1:
+        if i_split_head >= len_text:
+            break  # 结尾一定有标点，所以直接跳出即可，最后一段在上次已加入
+        if todo_text[i_split_head] in splits:
+            i_split_head += 1
+            todo_texts.append(todo_text[i_split_tail:i_split_head])
+            i_split_tail = i_split_head
+        else:
+            i_split_head += 1
+    return todo_texts
+
+
+# 不切
+@register_method("cut0")
+def cut0(inp):
+    if not set(inp).issubset(punctuation):
+        return inp
+    else:
+        return "/n"
+
+
+# 凑四句一切
+@register_method("cut1")
+def cut1(inp):
+    inp = inp.strip("\n")
+    inps = split(inp)
+    split_idx = list(range(0, len(inps), 4))
+    split_idx[-1] = None
+    if len(split_idx) > 1:
+        opts = []
+        for idx in range(len(split_idx) - 1):
+            opts.append("".join(inps[split_idx[idx]: split_idx[idx + 1]]))
+    else:
+        opts = [inp]
+    opts = [item for item in opts if not set(item).issubset(punctuation)]
+    return "\n".join(opts)
+
+
+# 凑50字一切
+@register_method("cut2")
+def cut2(inp):
+    inp = inp.strip("\n")
+    inps = split(inp)
+    if len(inps) < 2:
+        return inp
+    opts = []
+    summ = 0
+    tmp_str = ""
+    for i in range(len(inps)):
+        summ += len(inps[i])
+        tmp_str += inps[i]
+        if summ > 50:
+            summ = 0
+            opts.append(tmp_str)
+            tmp_str = ""
+    if tmp_str != "":
+        opts.append(tmp_str)
+    # print(opts)
+    if len(opts) > 1 and len(opts[-1]) < 50:  ##如果最后一个太短了，和前一个合一起
+        opts[-2] = opts[-2] + opts[-1]
+        opts = opts[:-1]
+    opts = [item for item in opts if not set(item).issubset(punctuation)]
+    return "\n".join(opts)
+
+# 按中文句号。切
+@register_method("cut3")
+def cut3(inp):
+    inp = inp.strip("\n")
+    opts = ["%s" % item for item in inp.strip("。").split("。")]
+    opts = [item for item in opts if not set(item).issubset(punctuation)]
+    return "\n".join(opts)
+
+#按英文句号.切
+@register_method("cut4")
+def cut4(inp):
+    inp = inp.strip("\n")
+    opts = ["%s" % item for item in inp.strip(".").split(".")]
+    opts = [item for item in opts if not set(item).issubset(punctuation)]
+    return "\n".join(opts)
+
+# 按标点符号切
+# contributed by https://github.com/AI-Hobbyist/GPT-SoVITS/blob/main/GPT_SoVITS/inference_webui.py
+@register_method("cut5")
+def cut5(inp):
+    inp = inp.strip("\n")
+    punds = {',', '.', ';', '?', '!', '、', '，', '。', '？', '！', ';', '：', '…'}
+    mergeitems = []
+    items = []
+
+    for i, char in enumerate(inp):
+        if char in punds:
+            if char == '.' and i > 0 and i < len(inp) - 1 and inp[i - 1].isdigit() and inp[i + 1].isdigit():
+                items.append(char)
+            else:
+                items.append(char)
+                mergeitems.append("".join(items))
+                items = []
+        else:
+            items.append(char)
+
+    if items:
+        mergeitems.append("".join(items))
+
+    opt = [item for item in mergeitems if not set(item).issubset(punds)]
+    return "\n".join(opt)
+
+
+
+if __name__ == '__main__':
+    method = get_method("cut5")
+    print(method("你好，我是小明。你好，我是小红。你好，我是小刚。你好，我是小张。"))
+    

TTS_infer_pack/tts_infer_module.py

@@ -0,0 +1,150 @@
+import os
+import random
+import typing
+from pathlib import Path
+
+from .TTS import TTS, TTS_Config
+
+# from pydub import AudioSegment
+# from pydub.playback import play
+# from setting import settings
+current_dir = Path(__file__).parent
+parent_dir = current_dir.parent
+
+
+class TTSModule:
+    def __init__(self, tts_config: TTS_Config):
+        self.tts_pipeline = TTS(tts_config)
+
+    def setup_inference_params(
+            self,
+            text_lang: typing.Literal["en", "all_zh", "all_ja"] = 'en',
+            ref_audio: str = '',
+            prompt_text: str = '',
+            prompt_lang: str = '',
+            top_k: int = 60,
+            top_p: float = 1.0,
+            temperature: float = 1.0,
+            text_split_method: typing.Literal['cut0', 'cut1', 'cut2', 'cut3', 'cut4', 'cut5'] = 'cut5',
+            batch_size: int = 100,
+            speed_factor: float = 1.1,
+            split_bucket: bool = True,
+            return_fragment: bool = False,
+            fragment_interval: float = 0.1,
+            seed: int = 233333,
+            parallel_infer: bool = True,
+    ):
+        """
+        Set up inference parameters for the model.
+
+        :param text_lang: Language of the input text. Choose from "en" (English), "all_zh" (Chinese), or "all_ja" (Japanese).
+        :param ref_audio: Path to the reference audio file. This audio is used as a voice reference for the output.
+        :param prompt_text: Text of the reference speech. This can be used to guide the model's output style or content.
+        :param prompt_lang: Language of the reference speech. Should correspond to the language of prompt_text.
+        :param top_k: Top-K sampling parameter for GPT. Limits the next token selection to the K most likely tokens.
+        :param top_p: Top-P (nucleus) sampling parameter for GPT. Sets a cumulative probability cutoff for token selection.
+        :param temperature: Temperature for GPT sampling. Higher values increase randomness, lower values make output more deterministic.
+        :param text_split_method: Method to split the input text:
+            - "cut0": No splitting (process whole text at once)
+            - "cut1": Split every 4 sentences
+            - "cut2": Split every 50 words
+            - "cut3": Split at Chinese full stops '。'
+            - "cut4": Split at English periods '.'
+            - "cut5": Automatic splitting (recommended)
+        :param batch_size: Batch size for inference. Larger values may speed up processing but require more memory.
+        :param speed_factor: Factor to control the speed of the output audio. Values > 1 speed up, < 1 slow down.
+        :param split_bucket: Whether to use bucket splitting for more efficient processing. Recommended to keep on.
+        :param return_fragment: If True, return individual fragments of the generated audio.
+        :param fragment_interval: Time interval (in seconds) between each sentence or fragment in the output audio.
+        :param seed: Random seed for reproducibility. Set a fixed value for consistent results across runs.
+        :param parallel_infer: whether to use parallel inference.
+        """
+        self.text_lang = text_lang
+        self.ref_audio = str(ref_audio)
+        self.prompt_text = prompt_text
+        self.prompt_lang = prompt_lang
+        self.top_k = max(1, int(top_k))  # 确保top_k至少为1
+        self.top_p = max(0.0, min(1.0, float(top_p)))  # 确保top_p在0到1之间
+        self.temperature = max(0.0, float(temperature))  # 确保temperature非负
+        self.text_split_method = text_split_method
+        self.batch_size = max(1, int(batch_size))  # 确保batch_size至少为1
+        self.speed_factor = max(0.1, float(speed_factor))  # 确保speed_factor至少为0.1
+        self.split_bucket = bool(split_bucket)
+        self.return_fragment = bool(return_fragment)
+        self.fragment_interval = max(0.0, float(fragment_interval))  # 确保fragment_interval非负
+        self.seed = int(seed)
+        self.parallel_infer = parallel_infer
+
+        # Validate that the reference audio file exists
+        if self.ref_audio and not os.path.exists(self.ref_audio):
+            raise FileNotFoundError(f"Reference audio file not found: {self.ref_audio}")
+
+    def inference(self, text, text_lang, ref_audio_path, prompt_text, prompt_lang, top_k,
+                  top_p, temperature, text_split_method, batch_size, speed_factor,
+                  ref_text_free, split_bucket, fragment_interval, seed, parallel_infer: bool):
+        actual_seed = seed if seed not in [-1, "", None] else random.randrange(1 << 32)
+        inputs = {
+            "text": text,
+            "text_lang": text_lang,
+            "ref_audio_path": ref_audio_path,
+            "prompt_text": prompt_text if not ref_text_free else "",
+            "prompt_lang": prompt_lang,
+            "top_k": top_k,
+            "top_p": top_p,
+            "temperature": temperature,
+            "text_split_method": text_split_method,
+            "batch_size": int(batch_size),
+            "speed_factor": float(speed_factor),
+            "split_bucket": split_bucket,
+            "return_fragment": False,
+            "fragment_interval": fragment_interval,
+            "seed": actual_seed,
+            "parallel_infer": parallel_infer
+        }
+        # print(inputs)
+        for item in self.tts_pipeline.run(inputs):
+            yield item, actual_seed
+
+    def generate_audio(self, text, warmup: bool = False):
+
+        [output] = self.inference(
+            text,
+            self.text_lang,
+            self.ref_audio,
+            self.prompt_text,
+            self.prompt_lang,
+            self.top_k,
+            self.top_p,
+            self.temperature,
+            self.text_split_method,
+            self.batch_size,
+            self.speed_factor,
+            self.return_fragment,
+            self.split_bucket,
+            self.fragment_interval,
+            self.seed,
+            self.parallel_infer,
+        )
+
+        if warmup:
+            return None
+
+        return output
+
+        # solution v1
+        # print(f"Audio generated path : '{parent_dir}'")
+        # # output_folder = "./output"  # 指定输出目录
+        # output_folder = str(settings.AUDIO_OUTPUT_FOLDER)
+        # if not os.path.exists(output_folder):
+        #     os.makedirs(output_folder)  # 如果目录不存在，则创建它
+        # now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S%f")  # 获取当前时间的时间戳
+        # file_name = os.path.join(output_folder, f'output_{now}_{count}.wav')
+        # sf.write(file_name, output[0][1], samplerate=output[0][0], subtype='PCM_16')
+        #
+        # print(f"Audio file saved to: {file_name}")
+        # return file_name
+
+# 使用示例
+# tts_module = TTSModule("./moyoyo_tts/configs/tts_infer.yaml")
+# text = "Well, you know what the say, Families are like fudge, mostly sweet, but sometimes nuts. My family is doing great, thanks for asking! My son is growing up to be a smart and handsome young man, just like his mom. He's currently working on his own talker show, which I'm sure will be even more hilarious than mine."
+# tts_module.generate_audio(text)

__init__.py

@@ -0,0 +1,9 @@
+from .TTS_infer_pack.TTS import TTS_Config
+from .TTS_infer_pack.tts_infer_module import TTSModule
+from .utils import HParams
+
+__all__ = (
+    HParams,
+    TTSModule,
+    TTS_Config
+)

configs/.gitignore

@@ -0,0 +1 @@
+*.yaml

configs/s2.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32a167e10a5664e54be6e78f178bcb497c3ae1ea9a5c41297812db3cae6fec8a
+size 1556

download.py

@@ -0,0 +1,5 @@
+import os, sys
+now_dir = os.getcwd()
+sys.path.insert(0, now_dir)
+from text.g2pw import G2PWPinyin
+g2pw = G2PWPinyin(model_dir="moyoyo_tts/text/G2PWModel",model_source="moyoyo_tts/pretrained_models/chinese-roberta-wwm-ext-large",v_to_u=False, neutral_tone_with_five=True)

export_torch_script.py

@@ -0,0 +1,831 @@
+# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/models/t2s_model.py
+# reference: https://github.com/lifeiteng/vall-e
+import argparse
+from typing import Optional
+from my_utils import load_audio
+import torch
+import torchaudio
+
+from torch import IntTensor, LongTensor, Tensor, nn
+from torch.nn import functional as F
+
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+from feature_extractor import cnhubert
+
+from AR.models.t2s_lightning_module import Text2SemanticLightningModule
+from module.models_onnx import SynthesizerTrn
+
+from inference_webui import get_phones_and_bert
+
+import os
+import soundfile
+
+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,
+}
+
+def get_raw_t2s_model(dict_s1) -> Text2SemanticLightningModule:
+    config = dict_s1["config"]
+    config["model"]["dropout"] = float(config["model"]["dropout"])
+    t2s_model = Text2SemanticLightningModule(config, "****", is_train=False)
+    t2s_model.load_state_dict(dict_s1["weight"])
+    t2s_model = t2s_model.eval()
+    return t2s_model
+
+@torch.jit.script
+def logits_to_probs(
+    logits,
+    previous_tokens: Optional[torch.Tensor] = None,
+    temperature: float = 1.0,
+    top_k: Optional[int] = None,
+    top_p: Optional[int] = None,
+    repetition_penalty: float = 1.0,
+):
+    # if previous_tokens is not None:
+    #     previous_tokens = previous_tokens.squeeze()
+    # print(logits.shape,previous_tokens.shape)
+    # pdb.set_trace()
+    if previous_tokens is not None and repetition_penalty != 1.0:
+        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=1, index=sorted_indices, src=sorted_indices_to_remove
+        )
+        logits = logits.masked_fill(indices_to_remove, -float("Inf"))
+
+    logits = logits / max(temperature, 1e-5)
+
+    if top_k is not None:
+        v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+        pivot = v[: , -1].unsqueeze(-1)
+        logits = torch.where(logits < pivot, -float("Inf"), logits)
+
+    probs = torch.nn.functional.softmax(logits, dim=-1)
+    return probs
+
+@torch.jit.script
+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)
+
+@torch.jit.script
+def sample(
+    logits,
+    previous_tokens,
+    temperature: float = 1.0,
+    top_k: Optional[int] = None,
+    top_p: Optional[int] = None,
+    repetition_penalty: float = 1.0,
+):
+    probs = logits_to_probs(
+        logits=logits, previous_tokens=previous_tokens, temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty
+    )
+    idx_next = multinomial_sample_one_no_sync(probs)
+    return idx_next, probs
+
+
+@torch.jit.script
+def spectrogram_torch(y:Tensor, n_fft:int, sampling_rate:int, hop_size:int, win_size:int, center:bool=False):
+    hann_window = torch.hann_window(win_size,device=y.device,dtype=y.dtype)
+    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")
+
+@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: int,
+            hidden_dim: int,
+            mlp: T2SMLP,
+            qkv_w,
+            qkv_b,
+            out_w,
+            out_b,
+            norm_w1,
+            norm_b1,
+            norm_eps1: float,
+            norm_w2,
+            norm_b2,
+            norm_eps2: float,
+    ):
+        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):
+        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)
+
+        attn = F.scaled_dot_product_attention(q, k, v, ~attn_mask)
+
+        attn = attn.permute(2, 0, 1, 3).reshape(batch_size*q_len, self.hidden_dim)
+        attn = attn.view(q_len, batch_size, self.hidden_dim).transpose(1, 0)
+        attn = F.linear(self.to_mask(attn, padding_mask), self.out_w, self.out_b)
+
+        if padding_mask is not None:
+            for i in range(batch_size):
+                # mask = padding_mask[i,:,0]
+                if self.false.device!= padding_mask.device:
+                    self.false = self.false.to(padding_mask.device)
+                idx = torch.where(padding_mask[i,:,0]==self.false)[0]
+                x_item = x[i,idx,:].unsqueeze(0)
+                attn_item = attn[i,idx,:].unsqueeze(0)
+                x_item = x_item + attn_item
+                x_item = F.layer_norm(
+                    x_item, [self.hidden_dim], self.norm_w1, self.norm_b1, self.norm_eps1
+                )
+                x_item = x_item + self.mlp.forward(x_item)
+                x_item = F.layer_norm(
+                    x_item,
+                    [self.hidden_dim],
+                    self.norm_w2,
+                    self.norm_b2,
+                    self.norm_eps2,
+                )
+                x[i,idx,:] = x_item.squeeze(0)
+            x = self.to_mask(x, padding_mask)
+        else:
+            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 decode_next_token(self, x:torch.Tensor, k_cache:torch.Tensor, v_cache:torch.Tensor):
+        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.permute(2, 0, 1, 3).reshape(batch_size*q_len, self.hidden_dim)
+        attn = attn.view(q_len, batch_size, self.hidden_dim).transpose(1, 0)
+        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
+
+@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):
+        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)
+            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: list[torch.Tensor],
+        v_cache: list[torch.Tensor]):
+        for i in range(self.num_blocks):
+            x, k_cache[i], v_cache[i] = self.blocks[i].decode_next_token(x, k_cache[i], v_cache[i])
+        return x, k_cache, v_cache
+
+class VitsModel(nn.Module):
+    def __init__(self, vits_path):
+        super().__init__()
+        # dict_s2 = torch.load(vits_path,map_location="cpu")
+        dict_s2 = torch.load(vits_path)
+        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, speed=1.0):
+        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, speed)[0, 0]
+
+class T2SModel(nn.Module):
+    def __init__(self,raw_t2s:Text2SemanticLightningModule):
+        super(T2SModel, self).__init__()
+        self.model_dim = raw_t2s.model.model_dim
+        self.embedding_dim = raw_t2s.model.embedding_dim
+        self.num_head = raw_t2s.model.num_head
+        self.num_layers = raw_t2s.model.num_layers
+        self.vocab_size = raw_t2s.model.vocab_size
+        self.phoneme_vocab_size = raw_t2s.model.phoneme_vocab_size
+        # self.p_dropout = float(raw_t2s.model.p_dropout)
+        self.EOS:int = int(raw_t2s.model.EOS)
+        self.norm_first = raw_t2s.model.norm_first
+        assert self.EOS == self.vocab_size - 1
+        self.hz = 50
+
+        self.bert_proj = raw_t2s.model.bert_proj
+        self.ar_text_embedding = raw_t2s.model.ar_text_embedding
+        self.ar_text_position = raw_t2s.model.ar_text_position
+        self.ar_audio_embedding = raw_t2s.model.ar_audio_embedding
+        self.ar_audio_position = raw_t2s.model.ar_audio_position
+
+        # self.t2s_transformer = T2STransformer(self.num_layers, blocks)
+        # self.t2s_transformer = raw_t2s.model.t2s_transformer
+
+        blocks = []
+        h = raw_t2s.model.h
+
+        for i in range(self.num_layers):
+            layer = h.layers[i]
+            t2smlp = T2SMLP(
+                layer.linear1.weight,
+                layer.linear1.bias,
+                layer.linear2.weight,
+                layer.linear2.bias
+            )
+
+            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)
+
+        # self.ar_predict_layer = nn.Linear(self.model_dim, self.vocab_size, bias=False)
+        self.ar_predict_layer = raw_t2s.model.ar_predict_layer
+        # self.loss_fct = nn.CrossEntropyLoss(reduction="sum")
+        self.max_sec = raw_t2s.config["data"]["max_sec"]
+        self.top_k = int(raw_t2s.config["inference"]["top_k"])
+        self.early_stop_num = torch.LongTensor([self.hz * self.max_sec])
+
+    def forward(self,prompts:LongTensor, ref_seq:LongTensor, text_seq:LongTensor, ref_bert:torch.Tensor, text_bert:torch.Tensor):
+        bert = torch.cat([ref_bert.T, text_bert.T], 1)
+        all_phoneme_ids = torch.cat([ref_seq, text_seq], 1)
+        bert = bert.unsqueeze(0)
+
+        x = self.ar_text_embedding(all_phoneme_ids)
+        x = x + self.bert_proj(bert.transpose(1, 2))
+        x:torch.Tensor = self.ar_text_position(x)
+
+        early_stop_num = self.early_stop_num
+
+
+        #[1,N,512] [1,N]
+        # y, k, v, y_emb, x_example = self.first_stage_decoder(x, prompts)
+        y = prompts
+        # x_example = x[:,:,0] * 0.0
+
+        x_len = x.shape[1]
+        x_attn_mask = torch.zeros((x_len, x_len), dtype=torch.bool)
+
+        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),  ###xx的纯0扩展到xx纯0+xy纯1，(x,x+y)
+            value=True,
+        )
+        y_attn_mask = F.pad(  ###yy的右上1扩展到左边xy的0,(y,x+y)
+            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)
+
+        idx = 0
+
+        xy_dec, k_cache, v_cache = self.t2s_transformer.process_prompt(xy_pos, xy_attn_mask, None)
+
+        logits = self.ar_predict_layer(xy_dec[:, -1])
+        logits = logits[:, :-1]
+        samples = sample(logits, y, top_k=self.top_k, top_p=1, repetition_penalty=1.35, temperature=1.0)[0]
+        y = torch.concat([y, samples], dim=1)
+        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)
+
+        stop = False
+        # for idx in range(1, 50):
+        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]
+            # y, k, v, y_emb, logits, samples = self.stage_decoder(y, k, v, y_emb, x_example)
+            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<11):###至少预测出10个token不然不给停止（0.4s）
+                logits = logits[:, :-1]
+
+            samples = sample(logits, y, top_k=self.top_k, top_p=1, repetition_penalty=1.35, temperature=1.0)[0]
+
+            y = torch.concat([y, samples], dim=1)
+
+            if early_stop_num != -1 and (y.shape[1] - prefix_len) > 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)
+                break
+
+            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)
+
+        y[0,-1] = 0
+
+        return y[:, -idx:].unsqueeze(0)
+
+bert_path = os.environ.get(
+    "bert_path", "moyoyo_tts/pretrained_models/chinese-roberta-wwm-ext-large"
+)
+cnhubert_base_path = "moyoyo_tts/pretrained_models/chinese-hubert-base"
+cnhubert.cnhubert_base_path = cnhubert_base_path
+
+@torch.jit.script
+def build_phone_level_feature(res:Tensor, word2ph:IntTensor):
+    phone_level_feature = []
+    for i in range(word2ph.shape[0]):
+        repeat_feature = res[i].repeat(word2ph[i].item(), 1)
+        phone_level_feature.append(repeat_feature)
+    phone_level_feature = torch.cat(phone_level_feature, dim=0)
+    # [sum(word2ph), 1024]
+    return phone_level_feature
+
+class MyBertModel(torch.nn.Module):
+    def __init__(self, bert_model):
+        super(MyBertModel, self).__init__()
+        self.bert = bert_model
+
+    def forward(self, input_ids:torch.Tensor, attention_mask:torch.Tensor, token_type_ids:torch.Tensor, word2ph:IntTensor):
+        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
+        # res = torch.cat(outputs["hidden_states"][-3:-2], -1)[0][1:-1]
+        res = torch.cat(outputs[1][-3:-2], -1)[0][1:-1]
+        return build_phone_level_feature(res, word2ph)
+
+class SSLModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.ssl = cnhubert.get_model().model
+
+    def forward(self, ref_audio_16k)-> torch.Tensor:
+        ssl_content = self.ssl(ref_audio_16k)["last_hidden_state"].transpose(1, 2)
+        return ssl_content
+
+class ExportSSLModel(torch.nn.Module):
+    def __init__(self,ssl:SSLModel):
+        super().__init__()
+        self.ssl = ssl
+
+    def forward(self, ref_audio:torch.Tensor):
+        return self.ssl(ref_audio)
+
+    @torch.jit.export
+    def resample(self,ref_audio:torch.Tensor,src_sr:int,dst_sr:int)->torch.Tensor:
+        audio = resamplex(ref_audio,src_sr,dst_sr).float()
+        return audio
+
+def export_bert(output_path):
+    tokenizer = AutoTokenizer.from_pretrained(bert_path)
+
+    text = "叹息声一声接着一声传出,木兰对着房门织布.听不见织布机织布的声音,只听见木兰在叹息.问木兰在想什么?问木兰在惦记什么?木兰答道,我也没有在想什么,也没有在惦记什么."
+    ref_bert_inputs = tokenizer(text, return_tensors="pt")
+    word2ph = []
+    for c in text:
+        if c in ['，','。','：','？',",",".","?"]:
+            word2ph.append(1)
+        else:
+            word2ph.append(2)
+    ref_bert_inputs['word2ph'] = torch.Tensor(word2ph).int()
+
+    bert_model = AutoModelForMaskedLM.from_pretrained(bert_path,output_hidden_states=True,torchscript=True)
+    my_bert_model = MyBertModel(bert_model)
+
+    ref_bert_inputs = {
+        'input_ids': ref_bert_inputs['input_ids'],
+        'attention_mask': ref_bert_inputs['attention_mask'],
+        'token_type_ids': ref_bert_inputs['token_type_ids'],
+        'word2ph': ref_bert_inputs['word2ph']
+    }
+
+    torch._dynamo.mark_dynamic(ref_bert_inputs['input_ids'], 1)
+    torch._dynamo.mark_dynamic(ref_bert_inputs['attention_mask'], 1)
+    torch._dynamo.mark_dynamic(ref_bert_inputs['token_type_ids'], 1)
+    torch._dynamo.mark_dynamic(ref_bert_inputs['word2ph'], 0)
+
+    my_bert_model = torch.jit.trace(my_bert_model,example_kwarg_inputs=ref_bert_inputs)
+    output_path = os.path.join(output_path, "bert_model.pt")
+    my_bert_model.save(output_path)
+    print('#### exported bert ####')
+
+def export(gpt_path, vits_path, ref_audio_path, ref_text, output_path, export_bert_and_ssl=False, device='cpu'):
+    if not os.path.exists(output_path):
+        os.makedirs(output_path)
+        print(f"目录已创建: {output_path}")
+    else:
+        print(f"目录已存在: {output_path}")
+
+    ref_audio = torch.tensor([load_audio(ref_audio_path, 16000)]).float()
+    ssl = SSLModel()
+    if export_bert_and_ssl:
+        s = ExportSSLModel(torch.jit.trace(ssl,example_inputs=(ref_audio)))
+        ssl_path = os.path.join(output_path, "ssl_model.pt")
+        torch.jit.script(s).save(ssl_path)
+        print('#### exported ssl ####')
+        export_bert(output_path)
+    else:
+        s = ExportSSLModel(ssl)
+
+    print(f"device: {device}")
+
+
+    ref_seq_id,ref_bert_T,ref_norm_text = get_phones_and_bert(ref_text,"all_zh",'v2')
+    ref_seq = torch.LongTensor([ref_seq_id]).to(device)
+    ref_bert = ref_bert_T.T.to(ref_seq.device)
+    text_seq_id,text_bert_T,norm_text = get_phones_and_bert("这是一条测试语音，说什么无所谓，只是给它一个例子","all_zh",'v2')
+    text_seq = torch.LongTensor([text_seq_id]).to(device)
+    text_bert = text_bert_T.T.to(text_seq.device)
+
+    ssl_content = ssl(ref_audio).to(device)
+
+     # vits_path = "SoVITS_weights_v2/xw_e8_s216.pth"
+    vits = VitsModel(vits_path).to(device)
+    vits.eval()
+
+    # gpt_path = "GPT_weights_v2/xw-e15.ckpt"
+    # dict_s1 = torch.load(gpt_path, map_location=device)
+    dict_s1 = torch.load(gpt_path)
+    raw_t2s = get_raw_t2s_model(dict_s1).to(device)
+    print('#### get_raw_t2s_model ####')
+    print(raw_t2s.config)
+    t2s_m = T2SModel(raw_t2s)
+    t2s_m.eval()
+    t2s = torch.jit.script(t2s_m).to(device)
+    print('#### script t2s_m ####')
+
+    print("vits.hps.data.sampling_rate:",vits.hps.data.sampling_rate)
+    gpt_sovits = GPT_SoVITS(t2s,vits).to(device)
+    gpt_sovits.eval()
+
+    ref_audio_sr = s.resample(ref_audio,16000,32000).to(device)
+
+    torch._dynamo.mark_dynamic(ssl_content, 2)
+    torch._dynamo.mark_dynamic(ref_audio_sr, 1)
+    torch._dynamo.mark_dynamic(ref_seq, 1)
+    torch._dynamo.mark_dynamic(text_seq, 1)
+    torch._dynamo.mark_dynamic(ref_bert, 0)
+    torch._dynamo.mark_dynamic(text_bert, 0)
+
+    with torch.no_grad():
+        gpt_sovits_export = torch.jit.trace(
+        gpt_sovits,
+        example_inputs=(
+            ssl_content,
+            ref_audio_sr,
+            ref_seq,
+            text_seq,
+            ref_bert,
+            text_bert))
+
+        gpt_sovits_path = os.path.join(output_path, "gpt_sovits_model.pt")
+        gpt_sovits_export.save(gpt_sovits_path)
+        print('#### exported gpt_sovits ####')
+
+@torch.jit.script
+def parse_audio(ref_audio):
+    ref_audio_16k = torchaudio.functional.resample(ref_audio,48000,16000).float()#.to(ref_audio.device)
+    ref_audio_sr = torchaudio.functional.resample(ref_audio,48000,32000).float()#.to(ref_audio.device)
+    return ref_audio_16k,ref_audio_sr
+
+@torch.jit.script
+def resamplex(ref_audio:torch.Tensor,src_sr:int,dst_sr:int)->torch.Tensor:
+    return torchaudio.functional.resample(ref_audio,src_sr,dst_sr).float()
+
+class GPT_SoVITS(nn.Module):
+    def __init__(self, t2s:T2SModel,vits:VitsModel):
+        super().__init__()
+        self.t2s = t2s
+        self.vits = vits
+
+    def forward(self, ssl_content:torch.Tensor, ref_audio_sr:torch.Tensor, ref_seq:Tensor, text_seq:Tensor, ref_bert:Tensor, text_bert:Tensor, speed=1.0):
+        codes = self.vits.vq_model.extract_latent(ssl_content)
+        prompt_semantic = codes[0, 0]
+        prompts = prompt_semantic.unsqueeze(0)
+
+        pred_semantic = self.t2s(prompts, ref_seq, text_seq, ref_bert, text_bert)
+        audio = self.vits(text_seq, pred_semantic, ref_audio_sr, speed)
+        return audio
+
+def test():
+    parser = argparse.ArgumentParser(description="GPT-SoVITS Command Line Tool")
+    parser.add_argument('--gpt_model', required=True, help="Path to the GPT model file")
+    parser.add_argument('--sovits_model', required=True, help="Path to the SoVITS model file")
+    parser.add_argument('--ref_audio', required=True, help="Path to the reference audio file")
+    parser.add_argument('--ref_text', required=True, help="Path to the reference text file")
+    parser.add_argument('--output_path', required=True, help="Path to the output directory")
+
+
+    args = parser.parse_args()
+    gpt_path = args.gpt_model
+    vits_path = args.sovits_model
+    ref_audio_path = args.ref_audio
+    ref_text = args.ref_text
+
+    tokenizer = AutoTokenizer.from_pretrained(bert_path)
+    # bert_model = AutoModelForMaskedLM.from_pretrained(bert_path,output_hidden_states=True,torchscript=True)
+    # bert = MyBertModel(bert_model)
+    my_bert = torch.jit.load("onnx/bert_model.pt",map_location='cuda')
+
+    # dict_s1 = torch.load(gpt_path, map_location="cuda")
+    # raw_t2s = get_raw_t2s_model(dict_s1)
+    # t2s = T2SModel(raw_t2s)
+    # t2s.eval()
+    # t2s = torch.jit.load("onnx/xw/t2s_model.pt",map_location='cuda')
+
+    # vits_path = "SoVITS_weights_v2/xw_e8_s216.pth"
+    # vits = VitsModel(vits_path)
+    # vits.eval()
+
+    # ssl = ExportSSLModel(SSLModel()).to('cuda')
+    # ssl.eval()
+    ssl = torch.jit.load("onnx/by/ssl_model.pt",map_location='cuda')
+
+    # gpt_sovits = moyoyo_tts(t2s,vits)
+    gpt_sovits = torch.jit.load("onnx/by/gpt_sovits_model.pt",map_location='cuda')
+
+    ref_seq_id,ref_bert_T,ref_norm_text = get_phones_and_bert(ref_text,"all_zh",'v2')
+    ref_seq = torch.LongTensor([ref_seq_id])
+    ref_bert = ref_bert_T.T.to(ref_seq.device)
+    # text_seq_id,text_bert_T,norm_text = get_phones_and_bert("昨天晚上看见征兵文书,知道君主在大规模征兵,那么多卷征兵文册,每一卷上都有父亲的名字.","all_zh",'v2')
+    text = "昨天晚上看见征兵文书,知道君主在大规模征兵,那么多卷征兵文册,每一卷上都有父亲的名字."
+
+    text_seq_id,text_bert_T,norm_text = get_phones_and_bert(text,"all_zh",'v2')
+
+    test_bert = tokenizer(text, return_tensors="pt")
+    word2ph = []
+    for c in text:
+        if c in ['，','。','：','？',"?",",","."]:
+            word2ph.append(1)
+        else:
+            word2ph.append(2)
+    test_bert['word2ph'] = torch.Tensor(word2ph).int()
+
+    test_bert = my_bert(
+        test_bert['input_ids'].to('cuda'),
+        test_bert['attention_mask'].to('cuda'),
+        test_bert['token_type_ids'].to('cuda'),
+        test_bert['word2ph'].to('cuda')
+    )
+
+    text_seq = torch.LongTensor([text_seq_id])
+    text_bert = text_bert_T.T.to(text_seq.device)
+
+    print('text_bert:',text_bert.shape,text_bert)
+    print('test_bert:',test_bert.shape,test_bert)
+    print(torch.allclose(text_bert.to('cuda'),test_bert))
+
+    print('text_seq:',text_seq.shape)
+    print('text_bert:',text_bert.shape,text_bert.type())
+
+    #[1,N]
+    ref_audio = torch.tensor([load_audio(ref_audio_path, 16000)]).float().to('cuda')
+    print('ref_audio:',ref_audio.shape)
+
+    ref_audio_sr = ssl.resample(ref_audio,16000,32000)
+    print('start ssl')
+    ssl_content = ssl(ref_audio)
+
+    print('start gpt_sovits:')
+    print('ssl_content:',ssl_content.shape)
+    print('ref_audio_sr:',ref_audio_sr.shape)
+    print('ref_seq:',ref_seq.shape)
+    ref_seq=ref_seq.to('cuda')
+    print('text_seq:',text_seq.shape)
+    text_seq=text_seq.to('cuda')
+    print('ref_bert:',ref_bert.shape)
+    ref_bert=ref_bert.to('cuda')
+    print('text_bert:',text_bert.shape)
+    text_bert=text_bert.to('cuda')
+
+    with torch.no_grad():
+        audio = gpt_sovits(ssl_content, ref_audio_sr, ref_seq, text_seq, ref_bert, test_bert)
+    print('start write wav')
+    soundfile.write("out.wav", audio.detach().cpu().numpy(), 32000)
+
+
+import text
+import json
+
+def export_symbel(version='v2'):
+    if version=='v1':
+        symbols = text._symbol_to_id_v1
+        with open(f"onnx/symbols_v1.json", "w") as file:
+            json.dump(symbols, file, indent=4)
+    else:
+        symbols = text._symbol_to_id_v2
+        with open(f"onnx/symbols_v2.json", "w") as file:
+            json.dump(symbols, file, indent=4)
+
+def main():
+    parser = argparse.ArgumentParser(description="GPT-SoVITS Command Line Tool")
+    parser.add_argument('--gpt_model', required=True, help="Path to the GPT model file")
+    parser.add_argument('--sovits_model', required=True, help="Path to the SoVITS model file")
+    parser.add_argument('--ref_audio', required=True, help="Path to the reference audio file")
+    parser.add_argument('--ref_text', required=True, help="Path to the reference text file")
+    parser.add_argument('--output_path', required=True, help="Path to the output directory")
+    parser.add_argument('--export_common_model', action='store_true', help="Export Bert and SSL model")
+    parser.add_argument('--device', help="Device to use")
+
+    args = parser.parse_args()
+    export(
+        gpt_path=args.gpt_model,
+        vits_path=args.sovits_model,
+        ref_audio_path=args.ref_audio,
+        ref_text=args.ref_text,
+        output_path=args.output_path,
+        device=args.device,
+        export_bert_and_ssl=args.export_common_model,
+    )
+
+import inference_webui
+if __name__ == "__main__":
+    inference_webui.is_half=False
+    inference_webui.dtype=torch.float32
+    main()
+    # test()

feature_extractor/__init__.py

@@ -0,0 +1,6 @@
+from . import cnhubert, whisper_enc
+
+content_module_map = {
+    'cnhubert': cnhubert,
+    'whisper': whisper_enc
+}

feature_extractor/cnhubert.py

@@ -0,0 +1,108 @@
+import os
+
+import torch
+from transformers import logging as tf_logging
+
+tf_logging.set_verbosity_error()
+
+import logging
+logging.getLogger("numba").setLevel(logging.WARNING)
+
+from transformers import (
+    Wav2Vec2FeatureExtractor,
+    HubertModel,
+)
+
+import moyoyo_tts.utils as utils
+import torch.nn as nn
+
+cnhubert_base_path = None
+
+
+class CNHubert(nn.Module):
+    def __init__(self, base_path:str=None):
+        super().__init__()
+        if base_path is None:
+            base_path = cnhubert_base_path
+        if os.path.exists(base_path):...
+        else:raise FileNotFoundError(base_path)
+        self.model = HubertModel.from_pretrained(base_path, local_files_only=True)
+        self.feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
+            base_path, local_files_only=True
+        )
+
+    def forward(self, x):
+        input_values = self.feature_extractor(
+            x, return_tensors="pt", sampling_rate=16000
+        ).input_values.to(x.device)
+        feats = self.model(input_values)["last_hidden_state"]
+        return feats
+
+
+# class CNHubertLarge(nn.Module):
+#     def __init__(self):
+#         super().__init__()
+#         self.model = HubertModel.from_pretrained("/data/docker/liujing04/gpt-vits/chinese-hubert-large")
+#         self.feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("/data/docker/liujing04/gpt-vits/chinese-hubert-large")
+#     def forward(self, x):
+#         input_values = self.feature_extractor(x, return_tensors="pt", sampling_rate=16000).input_values.to(x.device)
+#         feats = self.model(input_values)["last_hidden_state"]
+#         return feats
+#
+# class CVec(nn.Module):
+#     def __init__(self):
+#         super().__init__()
+#         self.model = HubertModel.from_pretrained("/data/docker/liujing04/vc-webui-big/hubert_base")
+#         self.feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("/data/docker/liujing04/vc-webui-big/hubert_base")
+#     def forward(self, x):
+#         input_values = self.feature_extractor(x, return_tensors="pt", sampling_rate=16000).input_values.to(x.device)
+#         feats = self.model(input_values)["last_hidden_state"]
+#         return feats
+#
+# class cnw2v2base(nn.Module):
+#     def __init__(self):
+#         super().__init__()
+#         self.model = Wav2Vec2Model.from_pretrained("/data/docker/liujing04/gpt-vits/chinese-wav2vec2-base")
+#         self.feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("/data/docker/liujing04/gpt-vits/chinese-wav2vec2-base")
+#     def forward(self, x):
+#         input_values = self.feature_extractor(x, return_tensors="pt", sampling_rate=16000).input_values.to(x.device)
+#         feats = self.model(input_values)["last_hidden_state"]
+#         return feats
+
+
+def get_model():
+    model = CNHubert()
+    model.eval()
+    return model
+
+
+# def get_large_model():
+#     model = CNHubertLarge()
+#     model.eval()
+#     return model
+#
+# def get_model_cvec():
+#     model = CVec()
+#     model.eval()
+#     return model
+#
+# def get_model_cnw2v2base():
+#     model = cnw2v2base()
+#     model.eval()
+#     return model
+
+
+def get_content(hmodel, wav_16k_tensor):
+    with torch.no_grad():
+        feats = hmodel(wav_16k_tensor)
+    return feats.transpose(1, 2)
+
+
+if __name__ == "__main__":
+    model = get_model()
+    src_path = "/Users/Shared/原音频2.wav"
+    wav_16k_tensor = utils.load_wav_to_torch_and_resample(src_path, 16000)
+    model = model
+    wav_16k_tensor = wav_16k_tensor
+    feats = get_content(model, wav_16k_tensor)
+    print(feats.shape)

feature_extractor/whisper_enc.py

@@ -0,0 +1,25 @@
+import torch
+
+
+def get_model():
+    import whisper
+
+    model = whisper.load_model("small", device="cpu")
+
+    return model.encoder
+
+
+def get_content(model=None, wav_16k_tensor=None):
+    from whisper import log_mel_spectrogram, pad_or_trim
+
+    dev = next(model.parameters()).device
+    mel = log_mel_spectrogram(wav_16k_tensor).to(dev)[:, :3000]
+    # if torch.cuda.is_available():
+    #     mel = mel.to(torch.float16)
+    feature_len = mel.shape[-1] // 2
+    assert mel.shape[-1] < 3000, "输入音频过长，只允许输入30以内音频"
+    with torch.no_grad():
+        feature = model(pad_or_trim(mel, 3000).unsqueeze(0))[
+            :1, :feature_len, :
+        ].transpose(1, 2)
+    return feature

inference_gui.py

@@ -0,0 +1,310 @@
+import os
+import sys
+from PyQt5.QtCore import QEvent
+from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QLineEdit, QPushButton, QTextEdit
+from PyQt5.QtWidgets import QGridLayout, QVBoxLayout, QWidget, QFileDialog, QStatusBar, QComboBox
+import soundfile as sf
+
+from tools.i18n.i18n import I18nAuto
+i18n = I18nAuto()
+
+from inference_webui import gpt_path, sovits_path, change_gpt_weights, change_sovits_weights, get_tts_wav
+
+
+class GPTSoVITSGUI(QMainWindow):
+    GPT_Path = gpt_path
+    SoVITS_Path = sovits_path
+
+    def __init__(self):
+        super().__init__()
+
+        self.setWindowTitle('GPT-SoVITS GUI')
+        self.setGeometry(800, 450, 950, 850)
+
+        self.setStyleSheet("""
+            QWidget {
+                background-color: #a3d3b1; 
+            }
+
+            QTabWidget::pane {
+                background-color: #a3d3b1;  
+            }
+
+            QTabWidget::tab-bar {
+                alignment: left;
+            }
+
+            QTabBar::tab {
+                background: #8da4bf; 
+                color: #ffffff;  
+                padding: 8px;
+            }
+
+            QTabBar::tab:selected {
+                background: #2a3f54; 
+            }
+
+            QLabel {
+                color: #000000;  
+            }
+
+            QPushButton {
+                background-color: #4CAF50; 
+                color: white;  
+                padding: 8px;
+                border: 1px solid #4CAF50;
+                border-radius: 4px;
+            }
+
+            QPushButton:hover {
+                background-color: #45a049;  
+                border: 1px solid #45a049;
+                box-shadow: 2px 2px 2px rgba(0, 0, 0, 0.1);
+            }
+        """)    
+
+        license_text = (
+        "本软件以MIT协议开源, 作者不对软件具备任何控制力, 使用软件者、传播软件导出的声音者自负全责. "
+        "如不认可该条款, 则不能使用或引用软件包内任何代码和文件. 详见根目录LICENSE.")
+        license_label = QLabel(license_text)
+        license_label.setWordWrap(True)
+
+        self.GPT_model_label = QLabel("选择GPT模型:")
+        self.GPT_model_input = QLineEdit()
+        self.GPT_model_input.setPlaceholderText("拖拽或选择文件")
+        self.GPT_model_input.setText(self.GPT_Path)
+        self.GPT_model_input.setReadOnly(True)
+        self.GPT_model_button = QPushButton("选择GPT模型文件")
+        self.GPT_model_button.clicked.connect(self.select_GPT_model)
+
+        self.SoVITS_model_label = QLabel("选择SoVITS模型:")
+        self.SoVITS_model_input = QLineEdit()
+        self.SoVITS_model_input.setPlaceholderText("拖拽或选择文件")
+        self.SoVITS_model_input.setText(self.SoVITS_Path)
+        self.SoVITS_model_input.setReadOnly(True)
+        self.SoVITS_model_button = QPushButton("选择SoVITS模型文件")
+        self.SoVITS_model_button.clicked.connect(self.select_SoVITS_model)
+
+        self.ref_audio_label = QLabel("上传参考音频:")
+        self.ref_audio_input = QLineEdit()
+        self.ref_audio_input.setPlaceholderText("拖拽或选择文件")
+        self.ref_audio_input.setReadOnly(True)
+        self.ref_audio_button = QPushButton("选择音频文件")
+        self.ref_audio_button.clicked.connect(self.select_ref_audio)
+
+        self.ref_text_label = QLabel("参考音频文本:")
+        self.ref_text_input = QLineEdit()
+        self.ref_text_input.setPlaceholderText("直接输入文字或上传文本")
+        self.ref_text_button = QPushButton("上传文本")
+        self.ref_text_button.clicked.connect(self.upload_ref_text)
+
+        self.ref_language_label = QLabel("参考音频语言:")
+        self.ref_language_combobox = QComboBox()
+        self.ref_language_combobox.addItems(["中文", "英文", "日文", "中英混合", "日英混合", "多语种混合"])
+        self.ref_language_combobox.setCurrentText("多语种混合")
+
+        self.target_text_label = QLabel("合成目标文本:")
+        self.target_text_input = QLineEdit()
+        self.target_text_input.setPlaceholderText("直接输入文字或上传文本")
+        self.target_text_button = QPushButton("上传文本")
+        self.target_text_button.clicked.connect(self.upload_target_text)
+
+        self.target_language_label = QLabel("合成音频语言:")
+        self.target_language_combobox = QComboBox()
+        self.target_language_combobox.addItems(["中文", "英文", "日文", "中英混合", "日英混合", "多语种混合"])
+        self.target_language_combobox.setCurrentText("多语种混合")
+
+        self.output_label = QLabel("输出音频路径:")
+        self.output_input = QLineEdit()
+        self.output_input.setPlaceholderText("拖拽或选择文件")
+        self.output_input.setReadOnly(True)
+        self.output_button = QPushButton("选择文件夹")
+        self.output_button.clicked.connect(self.select_output_path)
+
+        self.output_text = QTextEdit()
+        self.output_text.setReadOnly(True)
+
+        self.add_drag_drop_events([
+            self.GPT_model_input,
+            self.SoVITS_model_input,
+            self.ref_audio_input,
+            self.ref_text_input,
+            self.target_text_input,
+            self.output_input,
+        ])
+
+        self.synthesize_button = QPushButton("合成")
+        self.synthesize_button.clicked.connect(self.synthesize)
+
+        self.clear_output_button = QPushButton("清空输出")
+        self.clear_output_button.clicked.connect(self.clear_output)
+
+        self.status_bar = QStatusBar()
+
+        main_layout = QVBoxLayout()
+
+        input_layout = QGridLayout(self)
+        input_layout.setSpacing(10)
+
+        input_layout.addWidget(license_label, 0, 0, 1, 3)
+
+        input_layout.addWidget(self.GPT_model_label, 1, 0)
+        input_layout.addWidget(self.GPT_model_input, 2, 0, 1, 2)
+        input_layout.addWidget(self.GPT_model_button, 2, 2)
+
+        input_layout.addWidget(self.SoVITS_model_label, 3, 0)
+        input_layout.addWidget(self.SoVITS_model_input, 4, 0, 1, 2)
+        input_layout.addWidget(self.SoVITS_model_button, 4, 2)
+
+        input_layout.addWidget(self.ref_audio_label, 5, 0)
+        input_layout.addWidget(self.ref_audio_input, 6, 0, 1, 2)
+        input_layout.addWidget(self.ref_audio_button, 6, 2)
+
+        input_layout.addWidget(self.ref_language_label, 7, 0)
+        input_layout.addWidget(self.ref_language_combobox, 8, 0, 1, 1)
+        input_layout.addWidget(self.ref_text_label, 9, 0)
+        input_layout.addWidget(self.ref_text_input, 10, 0, 1, 2)
+        input_layout.addWidget(self.ref_text_button, 10, 2)
+
+        input_layout.addWidget(self.target_language_label, 11, 0)
+        input_layout.addWidget(self.target_language_combobox, 12, 0, 1, 1)
+        input_layout.addWidget(self.target_text_label, 13, 0)
+        input_layout.addWidget(self.target_text_input, 14, 0, 1, 2)
+        input_layout.addWidget(self.target_text_button, 14, 2)
+
+        input_layout.addWidget(self.output_label, 15, 0)
+        input_layout.addWidget(self.output_input, 16, 0, 1, 2)
+        input_layout.addWidget(self.output_button, 16, 2)
+
+        main_layout.addLayout(input_layout)
+
+        output_layout = QVBoxLayout()
+        output_layout.addWidget(self.output_text)
+        main_layout.addLayout(output_layout)
+
+        main_layout.addWidget(self.synthesize_button)
+
+        main_layout.addWidget(self.clear_output_button)
+
+        main_layout.addWidget(self.status_bar)
+
+        self.central_widget = QWidget()
+        self.central_widget.setLayout(main_layout)
+        self.setCentralWidget(self.central_widget)
+
+    def dragEnterEvent(self, event):
+        if event.mimeData().hasUrls():
+            event.acceptProposedAction()
+
+    def dropEvent(self, event):
+        if event.mimeData().hasUrls():
+            file_paths = [url.toLocalFile() for url in event.mimeData().urls()]
+            if len(file_paths) == 1:
+                self.update_ref_audio(file_paths[0])
+            else:
+                self.update_ref_audio(", ".join(file_paths))
+
+    def add_drag_drop_events(self, widgets):
+        for widget in widgets:
+            widget.setAcceptDrops(True)
+            widget.installEventFilter(self)
+
+    def eventFilter(self, obj, event):
+        if event.type() in (QEvent.DragEnter, QEvent.Drop):
+            mime_data = event.mimeData()
+            if mime_data.hasUrls():
+                event.acceptProposedAction()
+
+        return super().eventFilter(obj, event)
+
+    def select_GPT_model(self):
+        file_path, _ = QFileDialog.getOpenFileName(self, "选择GPT模型文件", "", "GPT Files (*.ckpt)")
+        if file_path:
+            self.GPT_model_input.setText(file_path)
+
+    def select_SoVITS_model(self):
+        file_path, _ = QFileDialog.getOpenFileName(self, "选择SoVITS模型文件", "", "SoVITS Files (*.pth)")
+        if file_path:
+            self.SoVITS_model_input.setText(file_path)
+
+    def select_ref_audio(self):
+        file_path, _ = QFileDialog.getOpenFileName(self, "选择参考音频文件", "", "Audio Files (*.wav *.mp3)")
+        if file_path:
+            self.update_ref_audio(file_path)
+
+    def upload_ref_text(self):
+        file_path, _ = QFileDialog.getOpenFileName(self, "选择文本文件", "", "Text Files (*.txt)")
+        if file_path:
+            with open(file_path, 'r', encoding='utf-8') as file:
+                content = file.read()
+                self.ref_text_input.setText(content)
+
+    def upload_target_text(self):
+        file_path, _ = QFileDialog.getOpenFileName(self, "选择文本文件", "", "Text Files (*.txt)")
+        if file_path:
+            with open(file_path, 'r', encoding='utf-8') as file:
+                content = file.read()
+                self.target_text_input.setText(content)
+
+    def select_output_path(self):
+        options = QFileDialog.Options()
+        options |= QFileDialog.DontUseNativeDialog
+        options |= QFileDialog.ShowDirsOnly
+
+        folder_dialog = QFileDialog()
+        folder_dialog.setOptions(options)
+        folder_dialog.setFileMode(QFileDialog.Directory)
+
+        if folder_dialog.exec_():
+            folder_path = folder_dialog.selectedFiles()[0]
+            self.output_input.setText(folder_path)
+
+    def update_ref_audio(self, file_path):
+        self.ref_audio_input.setText(file_path)
+
+    def clear_output(self):
+        self.output_text.clear()
+
+    def synthesize(self):
+        GPT_model_path = self.GPT_model_input.text()
+        SoVITS_model_path = self.SoVITS_model_input.text()
+        ref_audio_path = self.ref_audio_input.text()
+        language_combobox = self.ref_language_combobox.currentText()
+        language_combobox = i18n(language_combobox)
+        ref_text = self.ref_text_input.text()
+        target_language_combobox = self.target_language_combobox.currentText()
+        target_language_combobox = i18n(target_language_combobox)
+        target_text = self.target_text_input.text()
+        output_path = self.output_input.text()
+
+        if GPT_model_path != self.GPT_Path:
+            change_gpt_weights(gpt_path=GPT_model_path)
+            self.GPT_Path = GPT_model_path
+        if SoVITS_model_path != self.SoVITS_Path:
+            change_sovits_weights(sovits_path=SoVITS_model_path)
+            self.SoVITS_Path = SoVITS_model_path
+
+        synthesis_result = get_tts_wav(ref_wav_path=ref_audio_path, 
+                                       prompt_text=ref_text, 
+                                       prompt_language=language_combobox, 
+                                       text=target_text, 
+                                       text_language=target_language_combobox)
+
+        result_list = list(synthesis_result)
+
+        if result_list:
+            last_sampling_rate, last_audio_data = result_list[-1]
+            output_wav_path = os.path.join(output_path, "output.wav") 
+            sf.write(output_wav_path, last_audio_data, last_sampling_rate)
+
+            result = "Audio saved to " + output_wav_path
+
+        self.status_bar.showMessage("合成完成！输出路径：" + output_wav_path, 5000)
+        self.output_text.append("处理结果：\n" + result)
+
+
+if __name__ == '__main__':
+    app = QApplication(sys.argv)
+    mainWin = GPTSoVITSGUI()
+    mainWin.show()
+    sys.exit(app.exec_())

inference_webui.py

@@ -0,0 +1,778 @@
+'''
+按中英混合识别
+按日英混合识别
+多语种启动切分识别语种
+全部按中文识别
+全部按英文识别
+全部按日文识别
+'''
+import json
+import logging
+import os
+import re
+import sys
+import traceback
+
+
+logging.getLogger("markdown_it").setLevel(logging.ERROR)
+logging.getLogger("urllib3").setLevel(logging.ERROR)
+logging.getLogger("httpcore").setLevel(logging.ERROR)
+logging.getLogger("httpx").setLevel(logging.ERROR)
+logging.getLogger("asyncio").setLevel(logging.ERROR)
+logging.getLogger("charset_normalizer").setLevel(logging.ERROR)
+logging.getLogger("torchaudio._extension").setLevel(logging.ERROR)
+logging.getLogger("multipart.multipart").setLevel(logging.ERROR)
+import LangSegment.LangSegment as LangSegment
+import torch
+
+try:
+    import gradio.analytics as analytics
+    analytics.version_check = lambda:None
+except:...
+
+version=os.environ.get("version","v2")
+pretrained_sovits_name=["moyoyo_tts/pretrained_models/gsv-v2final-pretrained/s2G2333k.pth", "moyoyo_tts/pretrained_models/s2G488k.pth"]
+pretrained_gpt_name=["moyoyo_tts/pretrained_models/gsv-v2final-pretrained/s1bert25hz-5kh-longer-epoch=12-step=369668.ckpt", "moyoyo_tts/pretrained_models/s1bert25hz-2kh-longer-epoch=68e-step=50232.ckpt"]
+
+_ =[[],[]]
+for i in range(2):
+    if os.path.exists(pretrained_gpt_name[i]):
+        _[0].append(pretrained_gpt_name[i])
+    if os.path.exists(pretrained_sovits_name[i]):
+        _[-1].append(pretrained_sovits_name[i])
+pretrained_gpt_name,pretrained_sovits_name = _
+
+
+
+if os.path.exists(f"./weight.json"):
+    pass
+else:
+    with open(f"./weight.json", 'w', encoding="utf-8") as file:json.dump({'GPT':{},'SoVITS':{}},file)
+
+with open(f"./weight.json", 'r', encoding="utf-8") as file:
+    weight_data = file.read()
+    weight_data=json.loads(weight_data)
+    gpt_path = os.environ.get(
+        "gpt_path", weight_data.get('GPT',{}).get(version,pretrained_gpt_name))
+    sovits_path = os.environ.get(
+        "sovits_path", weight_data.get('SoVITS',{}).get(version,pretrained_sovits_name))
+    if isinstance(gpt_path,list):
+        gpt_path = gpt_path[0]
+    if isinstance(sovits_path,list):
+        sovits_path = sovits_path[0]
+
+# gpt_path = os.environ.get(
+#     "gpt_path", pretrained_gpt_name
+# )
+# sovits_path = os.environ.get("sovits_path", pretrained_sovits_name)
+cnhubert_base_path = os.environ.get(
+    "cnhubert_base_path", "moyoyo_tts/pretrained_models/chinese-hubert-base"
+)
+bert_path = os.environ.get(
+    "bert_path", "moyoyo_tts/pretrained_models/chinese-roberta-wwm-ext-large"
+)
+infer_ttswebui = os.environ.get("infer_ttswebui", 9872)
+infer_ttswebui = int(infer_ttswebui)
+is_share = os.environ.get("is_share", "False")
+is_share = eval(is_share)
+if "_CUDA_VISIBLE_DEVICES" in os.environ:
+    os.environ["CUDA_VISIBLE_DEVICES"] = os.environ["_CUDA_VISIBLE_DEVICES"]
+is_half = eval(os.environ.get("is_half", "True")) and torch.cuda.is_available()
+punctuation = set(['!', '?', '…', ',', '.', '-'," "])
+import gradio as gr
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+import numpy as np
+import librosa
+from feature_extractor import cnhubert
+
+cnhubert.cnhubert_base_path = cnhubert_base_path
+
+from module.models import SynthesizerTrn
+from AR.models.t2s_lightning_module import Text2SemanticLightningModule
+from text import cleaned_text_to_sequence
+from text.cleaner import clean_text
+from time import time as ttime
+from module.mel_processing import spectrogram_torch
+from tools.my_utils import load_audio
+from tools.i18n.i18n import I18nAuto, scan_language_list
+
+language=os.environ.get("language","Auto")
+language=sys.argv[-1] if sys.argv[-1] in scan_language_list() else language
+i18n = I18nAuto(language=language)
+
+# os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'  # 确保直接启动推理UI时也能够设置。
+
+if torch.cuda.is_available():
+    device = "cuda"
+else:
+    device = "cpu"
+
+dict_language_v1 = {
+    i18n("中文"): "all_zh",#全部按中文识别
+    i18n("英文"): "en",#全部按英文识别#######不变
+    i18n("日文"): "all_ja",#全部按日文识别
+    i18n("中英混合"): "zh",#按中英混合识别####不变
+    i18n("日英混合"): "ja",#按日英混合识别####不变
+    i18n("多语种混合"): "auto",#多语种启动切分识别语种
+}
+dict_language_v2 = {
+    i18n("中文"): "all_zh",#全部按中文识别
+    i18n("英文"): "en",#全部按英文识别#######不变
+    i18n("日文"): "all_ja",#全部按日文识别
+    i18n("粤语"): "all_yue",#全部按中文识别
+    i18n("韩文"): "all_ko",#全部按韩文识别
+    i18n("中英混合"): "zh",#按中英混合识别####不变
+    i18n("日英混合"): "ja",#按日英混合识别####不变
+    i18n("粤英混合"): "yue",#按粤英混合识别####不变
+    i18n("韩英混合"): "ko",#按韩英混合识别####不变
+    i18n("多语种混合"): "auto",#多语种启动切分识别语种
+    i18n("多语种混合(粤语)"): "auto_yue",#多语种启动切分识别语种
+}
+dict_language = dict_language_v1 if version =='v1' else dict_language_v2
+
+tokenizer = AutoTokenizer.from_pretrained(bert_path)
+bert_model = AutoModelForMaskedLM.from_pretrained(bert_path)
+if is_half == True:
+    bert_model = bert_model.half().to(device)
+else:
+    bert_model = bert_model.to(device)
+
+
+def get_bert_feature(text, word2ph):
+    with torch.no_grad():
+        inputs = tokenizer(text, return_tensors="pt")
+        for i in inputs:
+            inputs[i] = inputs[i].to(device)
+        res = bert_model(**inputs, output_hidden_states=True)
+        res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu()[1:-1]
+    assert len(word2ph) == len(text)
+    phone_level_feature = []
+    for i in range(len(word2ph)):
+        repeat_feature = res[i].repeat(word2ph[i], 1)
+        phone_level_feature.append(repeat_feature)
+    phone_level_feature = torch.cat(phone_level_feature, dim=0)
+    return phone_level_feature.T
+
+
+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")
+
+
+ssl_model = cnhubert.get_model()
+if is_half == True:
+    ssl_model = ssl_model.half().to(device)
+else:
+    ssl_model = ssl_model.to(device)
+
+
+def change_sovits_weights(sovits_path,prompt_language=None,text_language=None):
+    global vq_model, hps, version, dict_language
+    dict_s2 = torch.load(sovits_path, map_location="cpu")
+    hps = dict_s2["config"]
+    hps = DictToAttrRecursive(hps)
+    hps.model.semantic_frame_rate = "25hz"
+    if 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
+    )
+    if ("pretrained" not in sovits_path):
+        del vq_model.enc_q
+    if is_half == True:
+        vq_model = vq_model.half().to(device)
+    else:
+        vq_model = vq_model.to(device)
+    vq_model.eval()
+    print(vq_model.load_state_dict(dict_s2["weight"], strict=False))
+    dict_language = dict_language_v1 if version =='v1' else dict_language_v2
+    with open("./weight.json")as f:
+        data=f.read()
+        data=json.loads(data)
+        data["SoVITS"][version]=sovits_path
+    with open("./weight.json","w")as f:f.write(json.dumps(data))
+    if prompt_language is not None and text_language is not None:
+        if prompt_language in list(dict_language.keys()):
+            prompt_text_update, prompt_language_update = {'__type__':'update'},  {'__type__':'update', 'value':prompt_language}
+        else:
+            prompt_text_update = {'__type__':'update', 'value':''}
+            prompt_language_update = {'__type__':'update', 'value':i18n("中文")}
+        if text_language in list(dict_language.keys()):
+            text_update, text_language_update = {'__type__':'update'}, {'__type__':'update', 'value':text_language}
+        else:
+            text_update = {'__type__':'update', 'value':''}
+            text_language_update = {'__type__':'update', 'value':i18n("中文")}
+        return  {'__type__':'update', 'choices':list(dict_language.keys())}, {'__type__':'update', 'choices':list(dict_language.keys())}, prompt_text_update, prompt_language_update, text_update, text_language_update
+
+
+
+change_sovits_weights(sovits_path)
+
+
+def change_gpt_weights(gpt_path):
+    global hz, max_sec, t2s_model, config
+    hz = 50
+    dict_s1 = torch.load(gpt_path, map_location="cpu")
+    config = dict_s1["config"]
+    max_sec = config["data"]["max_sec"]
+    t2s_model = Text2SemanticLightningModule(config, "****", is_train=False)
+    t2s_model.load_state_dict(dict_s1["weight"])
+    if is_half == True:
+        t2s_model = t2s_model.half()
+    t2s_model = t2s_model.to(device)
+    t2s_model.eval()
+    total = sum([param.nelement() for param in t2s_model.parameters()])
+    print("Number of parameter: %.2fM" % (total / 1e6))
+    with open("./weight.json")as f:
+        data=f.read()
+        data=json.loads(data)
+        data["GPT"][version]=gpt_path
+    with open("./weight.json","w")as f:f.write(json.dumps(data))
+
+
+change_gpt_weights(gpt_path)
+
+
+def get_spepc(hps, filename):
+    audio = load_audio(filename, int(hps.data.sampling_rate))
+    audio = torch.FloatTensor(audio)
+    maxx=audio.abs().max()
+    if(maxx>1):audio/=min(2,maxx)
+    audio_norm = audio
+    audio_norm = audio_norm.unsqueeze(0)
+    spec = spectrogram_torch(
+        audio_norm,
+        hps.data.filter_length,
+        hps.data.sampling_rate,
+        hps.data.hop_length,
+        hps.data.win_length,
+        center=False,
+    )
+    return spec
+
+def clean_text_inf(text, language, version):
+    phones, word2ph, norm_text = clean_text(text, language, version)
+    phones = cleaned_text_to_sequence(phones, version)
+    return phones, word2ph, norm_text
+
+dtype=torch.float16 if is_half == True else torch.float32
+def get_bert_inf(phones, word2ph, norm_text, language):
+    language=language.replace("all_","")
+    if language == "zh":
+        bert = get_bert_feature(norm_text, word2ph).to(device)#.to(dtype)
+    else:
+        bert = torch.zeros(
+            (1024, len(phones)),
+            dtype=torch.float16 if is_half == True else torch.float32,
+        ).to(device)
+
+    return bert
+
+
+splits = {"，", "。", "？", "！", ",", ".", "?", "!", "~", ":", "：", "—", "…", }
+
+
+def get_first(text):
+    pattern = "[" + "".join(re.escape(sep) for sep in splits) + "]"
+    text = re.split(pattern, text)[0].strip()
+    return text
+
+from text import chinese
+def get_phones_and_bert(text,language,version,final=False):
+    if language in {"en", "all_zh", "all_ja", "all_ko", "all_yue"}:
+        language = language.replace("all_","")
+        if language == "en":
+            LangSegment.setfilters(["en"])
+            formattext = " ".join(tmp["text"] for tmp in LangSegment.getTexts(text))
+        else:
+            # 因无法区别中日韩文汉字,以用户输入为准
+            formattext = text
+        while "  " in formattext:
+            formattext = formattext.replace("  ", " ")
+        if language == "zh":
+            if re.search(r'[A-Za-z]', formattext):
+                formattext = re.sub(r'[a-z]', lambda x: x.group(0).upper(), formattext)
+                formattext = chinese.mix_text_normalize(formattext)
+                return get_phones_and_bert(formattext,"zh",version)
+            else:
+                phones, word2ph, norm_text = clean_text_inf(formattext, language, version)
+                bert = get_bert_feature(norm_text, word2ph).to(device)
+        elif language == "yue" and re.search(r'[A-Za-z]', formattext):
+                formattext = re.sub(r'[a-z]', lambda x: x.group(0).upper(), formattext)
+                formattext = chinese.mix_text_normalize(formattext)
+                return get_phones_and_bert(formattext,"yue",version)
+        else:
+            phones, word2ph, norm_text = clean_text_inf(formattext, language, version)
+            bert = torch.zeros(
+                (1024, len(phones)),
+                dtype=torch.float16 if is_half == True else torch.float32,
+            ).to(device)
+    elif language in {"zh", "ja", "ko", "yue", "auto", "auto_yue"}:
+        textlist=[]
+        langlist=[]
+        LangSegment.setfilters(["zh","ja","en","ko"])
+        if language == "auto":
+            for tmp in LangSegment.getTexts(text):
+                langlist.append(tmp["lang"])
+                textlist.append(tmp["text"])
+        elif language == "auto_yue":
+            for tmp in LangSegment.getTexts(text):
+                if tmp["lang"] == "zh":
+                    tmp["lang"] = "yue"
+                langlist.append(tmp["lang"])
+                textlist.append(tmp["text"])
+        else:
+            for tmp in LangSegment.getTexts(text):
+                if tmp["lang"] == "en":
+                    langlist.append(tmp["lang"])
+                else:
+                    # 因无法区别中日韩文汉字,以用户输入为准
+                    langlist.append(language)
+                textlist.append(tmp["text"])
+        print(textlist)
+        print(langlist)
+        phones_list = []
+        bert_list = []
+        norm_text_list = []
+        for i in range(len(textlist)):
+            lang = langlist[i]
+            phones, word2ph, norm_text = clean_text_inf(textlist[i], lang, version)
+            bert = get_bert_inf(phones, word2ph, norm_text, lang)
+            phones_list.append(phones)
+            norm_text_list.append(norm_text)
+            bert_list.append(bert)
+        bert = torch.cat(bert_list, dim=1)
+        phones = sum(phones_list, [])
+        norm_text = ''.join(norm_text_list)
+
+    if not final and len(phones) < 6:
+        return get_phones_and_bert("." + text,language,version,final=True)
+
+    return phones,bert.to(dtype),norm_text
+
+
+def merge_short_text_in_array(texts, threshold):
+    if (len(texts)) < 2:
+        return texts
+    result = []
+    text = ""
+    for ele in texts:
+        text += ele
+        if len(text) >= threshold:
+            result.append(text)
+            text = ""
+    if (len(text) > 0):
+        if len(result) == 0:
+            result.append(text)
+        else:
+            result[len(result) - 1] += text
+    return result
+
+##ref_wav_path+prompt_text+prompt_language+text(单个)+text_language+top_k+top_p+temperature
+# cache_tokens={}#暂未实现清理机制
+cache= {}
+def get_tts_wav(ref_wav_path, prompt_text, prompt_language, text, text_language, how_to_cut=i18n("不切"), top_k=20, top_p=0.6, temperature=0.6, ref_free
+    =False,speed=1,if_freeze=False,inp_refs=None):
+    global cache
+    if ref_wav_path:pass
+    else:gr.Warning(i18n('请上传参考音频'))
+    if text:pass
+    else:gr.Warning(i18n('请填入推理文本'))
+    t = []
+    if prompt_text is None or len(prompt_text) == 0:
+        ref_free = True
+    t0 = ttime()
+    prompt_language = dict_language[prompt_language]
+    text_language = dict_language[text_language]
+
+
+    if not ref_free:
+        prompt_text = prompt_text.strip("\n")
+        if (prompt_text[-1] not in splits): prompt_text += "。" if prompt_language != "en" else "."
+        print(i18n("实际输入的参考文本:"), prompt_text)
+    text = text.strip("\n")
+    # if (text[0] not in splits and len(get_first(text)) < 4): text = "。" + text if text_language != "en" else "." + text
+
+    print(i18n("实际输入的目标文本:"), text)
+    zero_wav = np.zeros(
+        int(hps.data.sampling_rate * 0.3),
+        dtype=np.float16 if is_half == True else np.float32,
+    )
+    if not ref_free:
+        with torch.no_grad():
+            wav16k, sr = librosa.load(ref_wav_path, sr=16000)
+            if (wav16k.shape[0] > 160000 or wav16k.shape[0] < 48000):
+                gr.Warning(i18n("参考音频在3~10秒范围外，请更换！"))
+                raise OSError(i18n("参考音频在3~10秒范围外，请更换！"))
+            wav16k = torch.from_numpy(wav16k)
+            zero_wav_torch = torch.from_numpy(zero_wav)
+            if is_half == True:
+                wav16k = wav16k.half().to(device)
+                zero_wav_torch = zero_wav_torch.half().to(device)
+            else:
+                wav16k = wav16k.to(device)
+                zero_wav_torch = zero_wav_torch.to(device)
+            wav16k = torch.cat([wav16k, zero_wav_torch])
+            ssl_content = ssl_model.model(wav16k.unsqueeze(0))[
+                "last_hidden_state"
+            ].transpose(
+                1, 2
+            )  # .float()
+            codes = vq_model.extract_latent(ssl_content)
+            prompt_semantic = codes[0, 0]
+            prompt = prompt_semantic.unsqueeze(0).to(device)
+
+    t1 = ttime()
+    t.append(t1-t0)
+
+    if (how_to_cut == i18n("凑四句一切")):
+        text = cut1(text)
+    elif (how_to_cut == i18n("凑50字一切")):
+        text = cut2(text)
+    elif (how_to_cut == i18n("按中文句号。切")):
+        text = cut3(text)
+    elif (how_to_cut == i18n("按英文句号.切")):
+        text = cut4(text)
+    elif (how_to_cut == i18n("按标点符号切")):
+        text = cut5(text)
+    while "\n\n" in text:
+        text = text.replace("\n\n", "\n")
+    print(i18n("实际输入的目标文本(切句后):"), text)
+    texts = text.split("\n")
+    texts = process_text(texts)
+    texts = merge_short_text_in_array(texts, 5)
+    audio_opt = []
+    if not ref_free:
+        phones1,bert1,norm_text1=get_phones_and_bert(prompt_text, prompt_language, version)
+
+    for i_text,text in enumerate(texts):
+        # 解决输入目标文本的空行导致报错的问题
+        if (len(text.strip()) == 0):
+            continue
+        if (text[-1] not in splits): text += "。" if text_language != "en" else "."
+        print(i18n("实际输入的目标文本(每句):"), text)
+        phones2,bert2,norm_text2=get_phones_and_bert(text, text_language, version)
+        print(i18n("前端处理后的文本(每句):"), norm_text2)
+        if not ref_free:
+            bert = torch.cat([bert1, bert2], 1)
+            all_phoneme_ids = torch.LongTensor(phones1+phones2).to(device).unsqueeze(0)
+        else:
+            bert = bert2
+            all_phoneme_ids = torch.LongTensor(phones2).to(device).unsqueeze(0)
+
+        bert = bert.to(device).unsqueeze(0)
+        all_phoneme_len = torch.tensor([all_phoneme_ids.shape[-1]]).to(device)
+
+        t2 = ttime()
+        # cache_key="%s-%s-%s-%s-%s-%s-%s-%s"%(ref_wav_path,prompt_text,prompt_language,text,text_language,top_k,top_p,temperature)
+        # print(cache.keys(),if_freeze)
+        if(i_text in cache and if_freeze==True):pred_semantic=cache[i_text]
+        else:
+            with torch.no_grad():
+                pred_semantic, idx = t2s_model.model.infer_panel(
+                    all_phoneme_ids,
+                    all_phoneme_len,
+                    None if ref_free else prompt,
+                    bert,
+                    # prompt_phone_len=ph_offset,
+                    top_k=top_k,
+                    top_p=top_p,
+                    temperature=temperature,
+                    early_stop_num=hz * max_sec,
+                )
+                pred_semantic = pred_semantic[:, -idx:].unsqueeze(0)
+                cache[i_text]=pred_semantic
+        t3 = ttime()
+        refers=[]
+        if(inp_refs):
+            for path in inp_refs:
+                try:
+                    refer = get_spepc(hps, path.name).to(dtype).to(device)
+                    refers.append(refer)
+                except:
+                    traceback.print_exc()
+        if(len(refers)==0):refers = [get_spepc(hps, ref_wav_path).to(dtype).to(device)]
+        audio = (vq_model.decode(pred_semantic, torch.LongTensor(phones2).to(device).unsqueeze(0), refers,speed=speed).detach().cpu().numpy()[0, 0])
+        max_audio=np.abs(audio).max()#简单防止16bit爆音
+        if max_audio>1:audio/=max_audio
+        audio_opt.append(audio)
+        audio_opt.append(zero_wav)
+        t4 = ttime()
+        t.extend([t2 - t1,t3 - t2, t4 - t3])
+        t1 = ttime()
+    print("%.3f\t%.3f\t%.3f\t%.3f" %
+           (t[0], sum(t[1::3]), sum(t[2::3]), sum(t[3::3]))
+           )
+    yield hps.data.sampling_rate, (np.concatenate(audio_opt, 0) * 32768).astype(
+        np.int16
+    )
+
+
+def split(todo_text):
+    todo_text = todo_text.replace("……", "。").replace("——", "，")
+    if todo_text[-1] not in splits:
+        todo_text += "。"
+    i_split_head = i_split_tail = 0
+    len_text = len(todo_text)
+    todo_texts = []
+    while 1:
+        if i_split_head >= len_text:
+            break  # 结尾一定有标点，所以直接跳出即可，最后一段在上次已加入
+        if todo_text[i_split_head] in splits:
+            i_split_head += 1
+            todo_texts.append(todo_text[i_split_tail:i_split_head])
+            i_split_tail = i_split_head
+        else:
+            i_split_head += 1
+    return todo_texts
+
+
+def cut1(inp):
+    inp = inp.strip("\n")
+    inps = split(inp)
+    split_idx = list(range(0, len(inps), 4))
+    split_idx[-1] = None
+    if len(split_idx) > 1:
+        opts = []
+        for idx in range(len(split_idx) - 1):
+            opts.append("".join(inps[split_idx[idx]: split_idx[idx + 1]]))
+    else:
+        opts = [inp]
+    opts = [item for item in opts if not set(item).issubset(punctuation)]
+    return "\n".join(opts)
+
+
+def cut2(inp):
+    inp = inp.strip("\n")
+    inps = split(inp)
+    if len(inps) < 2:
+        return inp
+    opts = []
+    summ = 0
+    tmp_str = ""
+    for i in range(len(inps)):
+        summ += len(inps[i])
+        tmp_str += inps[i]
+        if summ > 50:
+            summ = 0
+            opts.append(tmp_str)
+            tmp_str = ""
+    if tmp_str != "":
+        opts.append(tmp_str)
+    # print(opts)
+    if len(opts) > 1 and len(opts[-1]) < 50:  ##如果最后一个太短了，和前一个合一起
+        opts[-2] = opts[-2] + opts[-1]
+        opts = opts[:-1]
+    opts = [item for item in opts if not set(item).issubset(punctuation)]
+    return "\n".join(opts)
+
+
+def cut3(inp):
+    inp = inp.strip("\n")
+    opts = ["%s" % item for item in inp.strip("。").split("。")]
+    opts = [item for item in opts if not set(item).issubset(punctuation)]
+    return  "\n".join(opts)
+
+def cut4(inp):
+    inp = inp.strip("\n")
+    opts = ["%s" % item for item in inp.strip(".").split(".")]
+    opts = [item for item in opts if not set(item).issubset(punctuation)]
+    return "\n".join(opts)
+
+
+# contributed by https://github.com/AI-Hobbyist/GPT-SoVITS/blob/main/GPT_SoVITS/inference_webui.py
+def cut5(inp):
+    inp = inp.strip("\n")
+    punds = {',', '.', ';', '?', '!', '、', '，', '。', '？', '！', ';', '：', '…'}
+    mergeitems = []
+    items = []
+
+    for i, char in enumerate(inp):
+        if char in punds:
+            if char == '.' and i > 0 and i < len(inp) - 1 and inp[i - 1].isdigit() and inp[i + 1].isdigit():
+                items.append(char)
+            else:
+                items.append(char)
+                mergeitems.append("".join(items))
+                items = []
+        else:
+            items.append(char)
+
+    if items:
+        mergeitems.append("".join(items))
+
+    opt = [item for item in mergeitems if not set(item).issubset(punds)]
+    return "\n".join(opt)
+
+
+def custom_sort_key(s):
+    # 使用正则表达式提取字符串中的数字部分和非数字部分
+    parts = re.split('(\d+)', s)
+    # 将数字部分转换为整数，非数字部分保持不变
+    parts = [int(part) if part.isdigit() else part for part in parts]
+    return parts
+
+def process_text(texts):
+    _text=[]
+    if all(text in [None, " ", "\n",""] for text in texts):
+        raise ValueError(i18n("请输入有效文本"))
+    for text in texts:
+        if text in  [None, " ", ""]:
+            pass
+        else:
+            _text.append(text)
+    return _text
+
+
+def change_choices():
+    SoVITS_names, GPT_names = get_weights_names(GPT_weight_root, SoVITS_weight_root)
+    return {"choices": sorted(SoVITS_names, key=custom_sort_key), "__type__": "update"}, {"choices": sorted(GPT_names, key=custom_sort_key), "__type__": "update"}
+
+
+# SoVITS_weight_root=["SoVITS_weights_v2","SoVITS_weights"]
+# GPT_weight_root=["GPT_weights_v2","GPT_weights"]
+SoVITS_weight_root=["checkpoints"]
+GPT_weight_root=["checkpoints"]
+for path in SoVITS_weight_root+GPT_weight_root:
+    os.makedirs(path,exist_ok=True)
+
+
+def get_weights_names(GPT_weight_root, SoVITS_weight_root):
+    SoVITS_names = [i for i in pretrained_sovits_name]
+    for path in SoVITS_weight_root:
+        for name in os.listdir(path):
+            if name.endswith(".pth"): SoVITS_names.append("%s/%s" % (path, name))
+    GPT_names = [i for i in pretrained_gpt_name]
+    for path in GPT_weight_root:
+        for name in os.listdir(path):
+            if name.endswith(".ckpt"): GPT_names.append("%s/%s" % (path, name))
+    return SoVITS_names, GPT_names
+
+
+SoVITS_names, GPT_names = get_weights_names(GPT_weight_root, SoVITS_weight_root)
+
+def html_center(text, label='p'):
+    return f"""<div style="text-align: center; margin: 100; padding: 50;">
+                <{label} style="margin: 0; padding: 0;">{text}</{label}>
+                </div>"""
+
+def html_left(text, label='p'):
+    return f"""<div style="text-align: left; margin: 0; padding: 0;">
+                <{label} style="margin: 0; padding: 0;">{text}</{label}>
+                </div>"""
+
+
+with gr.Blocks(title="GPT-SoVITS WebUI") as app:
+    gr.Markdown(
+        value=i18n("本软件以MIT协议开源, 作者不对软件具备任何控制力, 使用软件者、传播软件导出的声音者自负全责. <br>如不认可该条款, 则不能使用或引用软件包内任何代码和文件. 详见根目录<b>LICENSE</b>.")
+    )
+    with gr.Group():
+        gr.Markdown(html_center(i18n("模型切换"),'h3'))
+        with gr.Row():
+            GPT_dropdown = gr.Dropdown(label=i18n("GPT模型列表"), choices=sorted(GPT_names, key=custom_sort_key), value=gpt_path, interactive=True, scale=14)
+            SoVITS_dropdown = gr.Dropdown(label=i18n("SoVITS模型列表"), choices=sorted(SoVITS_names, key=custom_sort_key), value=sovits_path, interactive=True, scale=14)
+            refresh_button = gr.Button(i18n("刷新模型路径"), variant="primary", scale=14)
+            refresh_button.click(fn=change_choices, inputs=[], outputs=[SoVITS_dropdown, GPT_dropdown])
+        gr.Markdown(html_center(i18n("*请上传并填写参考信息"),'h3'))
+        with gr.Row():
+            inp_ref = gr.Audio(label=i18n("请上传3~10秒内参考音频，超过会报错！"), type="filepath", scale=13)
+            with gr.Column(scale=13):
+                ref_text_free = gr.Checkbox(label=i18n("开启无参考文本模式。不填参考文本亦相当于开启。"), value=False, interactive=True, show_label=True,scale=1)
+                gr.Markdown(html_left(i18n("使用无参考文本模式时建议使用微调的GPT，听不清参考音频说的啥(不晓得写啥)可以开。<br>开启后无视填写的参考文本。")))
+                prompt_text = gr.Textbox(label=i18n("参考音频的文本"), value="", lines=5, max_lines=5,scale=1)
+            with gr.Column(scale=14):
+                prompt_language = gr.Dropdown(
+                    label=i18n("参考音频的语种"), choices=list(dict_language.keys()), value=i18n("中文"),
+                )
+                inp_refs = gr.File(label=i18n("可选项：通过拖拽多个文件上传多个参考音频（建议同性），平均融合他们的音色。如不填写此项，音色由左侧单个参考音频控制。如是微调模型，建议参考音频全部在微调训练集音色内，底模不用管。"),file_count="multiple")
+        gr.Markdown(html_center(i18n("*请填写需要合成的目标文本和语种模式"),'h3'))
+        with gr.Row():
+            with gr.Column(scale=13):
+                text = gr.Textbox(label=i18n("需要合成的文本"), value="", lines=26, max_lines=26)
+            with gr.Column(scale=7):
+                text_language = gr.Dropdown(
+                        label=i18n("需要合成的语种")+i18n(".限制范围越小判别效果越好。"), choices=list(dict_language.keys()), value=i18n("中文"), scale=1
+                    )
+                how_to_cut = gr.Dropdown(
+                        label=i18n("怎么切"),
+                        choices=[i18n("不切"), i18n("凑四句一切"), i18n("凑50字一切"), i18n("按中文句号。切"), i18n("按英文句号.切"), i18n("按标点符号切"), ],
+                        value=i18n("凑四句一切"),
+                        interactive=True, scale=1
+                    )
+                gr.Markdown(value=html_center(i18n("语速调整，高为更快")))
+                if_freeze=gr.Checkbox(label=i18n("是否直接对上次合成结果调整语速和音色。防止随机性。"), value=False, interactive=True,show_label=True, scale=1)
+                speed = gr.Slider(minimum=0.6,maximum=1.65,step=0.05,label=i18n("语速"),value=1,interactive=True, scale=1)
+                gr.Markdown(html_center(i18n("GPT采样参数(无参考文本时不要太低。不懂就用默认)：")))
+                top_k = gr.Slider(minimum=1,maximum=100,step=1,label=i18n("top_k"),value=15,interactive=True, scale=1)
+                top_p = gr.Slider(minimum=0,maximum=1,step=0.05,label=i18n("top_p"),value=1,interactive=True, scale=1)
+                temperature = gr.Slider(minimum=0,maximum=1,step=0.05,label=i18n("temperature"),value=1,interactive=True,  scale=1)
+            # with gr.Column():
+            #     gr.Markdown(value=i18n("手工调整音素。当音素框不为空时使用手工音素输入推理，无视目标文本框。"))
+            #     phoneme=gr.Textbox(label=i18n("音素框"), value="")
+            #     get_phoneme_button = gr.Button(i18n("目标文本转音素"), variant="primary")
+        with gr.Row():
+            inference_button = gr.Button(i18n("合成语音"), variant="primary", size='lg', scale=25)
+            output = gr.Audio(label=i18n("输出的语音"), scale=14)
+
+        inference_button.click(
+            get_tts_wav,
+            [inp_ref, prompt_text, prompt_language, text, text_language, how_to_cut, top_k, top_p, temperature, ref_text_free,speed,if_freeze,inp_refs],
+            [output],
+        )
+        SoVITS_dropdown.change(change_sovits_weights, [SoVITS_dropdown,prompt_language,text_language], [prompt_language,text_language,prompt_text,prompt_language,text,text_language])
+        GPT_dropdown.change(change_gpt_weights, [GPT_dropdown], [])
+
+        # gr.Markdown(value=i18n("文本切分工具。太长的文本合成出来效果不一定好，所以太长建议先切。合成会根据文本的换行分开合成再拼起来。"))
+        # with gr.Row():
+        #     text_inp = gr.Textbox(label=i18n("需要合成的切分前文本"), value="")
+        #     button1 = gr.Button(i18n("凑四句一切"), variant="primary")
+        #     button2 = gr.Button(i18n("凑50字一切"), variant="primary")
+        #     button3 = gr.Button(i18n("按中文句号。切"), variant="primary")
+        #     button4 = gr.Button(i18n("按英文句号.切"), variant="primary")
+        #     button5 = gr.Button(i18n("按标点符号切"), variant="primary")
+        #     text_opt = gr.Textbox(label=i18n("切分后文本"), value="")
+        #     button1.click(cut1, [text_inp], [text_opt])
+        #     button2.click(cut2, [text_inp], [text_opt])
+        #     button3.click(cut3, [text_inp], [text_opt])
+        #     button4.click(cut4, [text_inp], [text_opt])
+        #     button5.click(cut5, [text_inp], [text_opt])
+        # gr.Markdown(html_center(i18n("后续将支持转音素、手工修改音素、语音合成分步执行。")))
+
+if __name__ == '__main__':
+    app.queue().launch(#concurrency_count=511, max_size=1022
+        server_name="0.0.0.0",
+        inbrowser=True,
+        share=is_share,
+        server_port=infer_ttswebui,
+        quiet=True,
+    )

inference_webui_fast.py

@@ -0,0 +1,336 @@
+'''
+按中英混合识别
+按日英混合识别
+多语种启动切分识别语种
+全部按中文识别
+全部按英文识别
+全部按日文识别
+'''
+import random
+import os, re, logging
+import sys
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+sys.path.append("%s/moyoyo_tts" % (now_dir))
+
+logging.getLogger("markdown_it").setLevel(logging.ERROR)
+logging.getLogger("urllib3").setLevel(logging.ERROR)
+logging.getLogger("httpcore").setLevel(logging.ERROR)
+logging.getLogger("httpx").setLevel(logging.ERROR)
+logging.getLogger("asyncio").setLevel(logging.ERROR)
+logging.getLogger("charset_normalizer").setLevel(logging.ERROR)
+logging.getLogger("torchaudio._extension").setLevel(logging.ERROR)
+import pdb
+import torch
+
+try:
+    import gradio.analytics as analytics
+    analytics.version_check = lambda:None
+except:...
+
+
+infer_ttswebui = os.environ.get("infer_ttswebui", 9872)
+infer_ttswebui = int(infer_ttswebui)
+is_share = os.environ.get("is_share", "False")
+is_share = eval(is_share)
+if "_CUDA_VISIBLE_DEVICES" in os.environ:
+    os.environ["CUDA_VISIBLE_DEVICES"] = os.environ["_CUDA_VISIBLE_DEVICES"]
+
+is_half = eval(os.environ.get("is_half", "True")) and torch.cuda.is_available()
+gpt_path = os.environ.get("gpt_path", None)
+sovits_path = os.environ.get("sovits_path", None)
+cnhubert_base_path = os.environ.get("cnhubert_base_path", None)
+bert_path = os.environ.get("bert_path", None)
+version=os.environ.get("version","v2")
+
+import gradio as gr
+from TTS_infer_pack.TTS import TTS, TTS_Config
+from TTS_infer_pack.text_segmentation_method import get_method
+from tools.i18n.i18n import I18nAuto, scan_language_list
+
+language=os.environ.get("language","Auto")
+language=sys.argv[-1] if sys.argv[-1] in scan_language_list() else language
+i18n = I18nAuto(language=language)
+
+
+# os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'  # 确保直接启动推理UI时也能够设置。
+
+if torch.cuda.is_available():
+    device = "cuda"
+# elif torch.backends.mps.is_available():
+#     device = "mps"
+else:
+    device = "cpu"
+
+dict_language_v1 = {
+    i18n("中文"): "all_zh",#全部按中文识别
+    i18n("英文"): "en",#全部按英文识别#######不变
+    i18n("日文"): "all_ja",#全部按日文识别
+    i18n("中英混合"): "zh",#按中英混合识别####不变
+    i18n("日英混合"): "ja",#按日英混合识别####不变
+    i18n("多语种混合"): "auto",#多语种启动切分识别语种
+}
+dict_language_v2 = {
+    i18n("中文"): "all_zh",#全部按中文识别
+    i18n("英文"): "en",#全部按英文识别#######不变
+    i18n("日文"): "all_ja",#全部按日文识别
+    i18n("粤语"): "all_yue",#全部按中文识别
+    i18n("韩文"): "all_ko",#全部按韩文识别
+    i18n("中英混合"): "zh",#按中英混合识别####不变
+    i18n("日英混合"): "ja",#按日英混合识别####不变
+    i18n("粤英混合"): "yue",#按粤英混合识别####不变
+    i18n("韩英混合"): "ko",#按韩英混合识别####不变
+    i18n("多语种混合"): "auto",#多语种启动切分识别语种
+    i18n("多语种混合(粤语)"): "auto_yue",#多语种启动切分识别语种
+}
+dict_language = dict_language_v1 if version =='v1' else dict_language_v2
+
+cut_method = {
+    i18n("不切"):"cut0",
+    i18n("凑四句一切"): "cut1",
+    i18n("凑50字一切"): "cut2",
+    i18n("按中文句号。切"): "cut3",
+    i18n("按英文句号.切"): "cut4",
+    i18n("按标点符号切"): "cut5",
+}
+
+tts_config = TTS_Config("moyoyo_tts/configs/tts_infer.yaml")
+tts_config.device = device
+tts_config.is_half = is_half
+tts_config.version = version
+if gpt_path is not None:
+    tts_config.t2s_weights_path = gpt_path
+if sovits_path is not None:
+    tts_config.vits_weights_path = sovits_path
+if cnhubert_base_path is not None:
+    tts_config.cnhuhbert_base_path = cnhubert_base_path
+if bert_path is not None:
+    tts_config.bert_base_path = bert_path
+
+print(tts_config)
+tts_pipeline = TTS(tts_config)
+gpt_path = tts_config.t2s_weights_path
+sovits_path = tts_config.vits_weights_path
+version = tts_config.version
+
+def inference(text, text_lang,
+              ref_audio_path,
+              aux_ref_audio_paths,
+              prompt_text,
+              prompt_lang, top_k,
+              top_p, temperature,
+              text_split_method, batch_size,
+              speed_factor, ref_text_free,
+              split_bucket,fragment_interval,
+              seed, keep_random, parallel_infer,
+              repetition_penalty
+              ):
+
+    seed = -1 if keep_random else seed
+    actual_seed = seed if seed not in [-1, "", None] else random.randrange(1 << 32)
+    inputs={
+        "text": text,
+        "text_lang": dict_language[text_lang],
+        "ref_audio_path": ref_audio_path,
+        "aux_ref_audio_paths": [item.name for item in aux_ref_audio_paths] if aux_ref_audio_paths is not None else [],
+        "prompt_text": prompt_text if not ref_text_free else "",
+        "prompt_lang": dict_language[prompt_lang],
+        "top_k": top_k,
+        "top_p": top_p,
+        "temperature": temperature,
+        "text_split_method": cut_method[text_split_method],
+        "batch_size":int(batch_size),
+        "speed_factor":float(speed_factor),
+        "split_bucket":split_bucket,
+        "return_fragment":False,
+        "fragment_interval":fragment_interval,
+        "seed":actual_seed,
+        "parallel_infer": parallel_infer,
+        "repetition_penalty": repetition_penalty,
+    }
+    for item in tts_pipeline.run(inputs):
+        yield item, actual_seed
+
+def custom_sort_key(s):
+    # 使用正则表达式提取字符串中的数字部分和非数字部分
+    parts = re.split('(\d+)', s)
+    # 将数字部分转换为整数，非数字部分保持不变
+    parts = [int(part) if part.isdigit() else part for part in parts]
+    return parts
+
+
+def change_choices():
+    SoVITS_names, GPT_names = get_weights_names(GPT_weight_root, SoVITS_weight_root)
+    return {"choices": sorted(SoVITS_names, key=custom_sort_key), "__type__": "update"}, {"choices": sorted(GPT_names, key=custom_sort_key), "__type__": "update"}
+
+
+pretrained_sovits_name=["moyoyo_tts/pretrained_models/gsv-v2final-pretrained/s2G2333k.pth", "moyoyo_tts/pretrained_models/s2G488k.pth"]
+pretrained_gpt_name=["moyoyo_tts/pretrained_models/gsv-v2final-pretrained/s1bert25hz-5kh-longer-epoch=12-step=369668.ckpt", "moyoyo_tts/pretrained_models/s1bert25hz-2kh-longer-epoch=68e-step=50232.ckpt"]
+_ =[[],[]]
+for i in range(2):
+    if os.path.exists(pretrained_gpt_name[i]):
+        _[0].append(pretrained_gpt_name[i])
+    if os.path.exists(pretrained_sovits_name[i]):
+        _[-1].append(pretrained_sovits_name[i])
+pretrained_gpt_name,pretrained_sovits_name = _
+
+SoVITS_weight_root=["SoVITS_weights_v2","SoVITS_weights"]
+GPT_weight_root=["GPT_weights_v2","GPT_weights"]
+for path in SoVITS_weight_root+GPT_weight_root:
+    os.makedirs(path,exist_ok=True)
+
+def get_weights_names(GPT_weight_root, SoVITS_weight_root):
+    SoVITS_names = [i for i in pretrained_sovits_name]
+    for path in SoVITS_weight_root:
+        for name in os.listdir(path):
+            if name.endswith(".pth"): SoVITS_names.append("%s/%s" % (path, name))
+    GPT_names = [i for i in pretrained_gpt_name]
+    for path in GPT_weight_root:
+        for name in os.listdir(path):
+            if name.endswith(".ckpt"): GPT_names.append("%s/%s" % (path, name))
+    return SoVITS_names, GPT_names
+
+
+SoVITS_names, GPT_names = get_weights_names(GPT_weight_root, SoVITS_weight_root)
+
+
+
+def change_sovits_weights(sovits_path,prompt_language=None,text_language=None):
+    tts_pipeline.init_vits_weights(sovits_path)
+    global version, dict_language
+    dict_language = dict_language_v1 if tts_pipeline.configs.version =='v1' else dict_language_v2
+    if prompt_language is not None and text_language is not None:
+        if prompt_language in list(dict_language.keys()):
+            prompt_text_update, prompt_language_update = {'__type__':'update'},  {'__type__':'update', 'value':prompt_language}
+        else:
+            prompt_text_update = {'__type__':'update', 'value':''}
+            prompt_language_update = {'__type__':'update', 'value':i18n("中文")}
+        if text_language in list(dict_language.keys()):
+            text_update, text_language_update = {'__type__':'update'}, {'__type__':'update', 'value':text_language}
+        else:
+            text_update = {'__type__':'update', 'value':''}
+            text_language_update = {'__type__':'update', 'value':i18n("中文")}
+        return  {'__type__':'update', 'choices':list(dict_language.keys())}, {'__type__':'update', 'choices':list(dict_language.keys())}, prompt_text_update, prompt_language_update, text_update, text_language_update
+
+
+
+with gr.Blocks(title="GPT-SoVITS WebUI") as app:
+    gr.Markdown(
+        value=i18n("本软件以MIT协议开源, 作者不对软件具备任何控制力, 使用软件者、传播软件导出的声音者自负全责. <br>如不认可该条款, 则不能使用或引用软件包内任何代码和文件. 详见根目录<b>LICENSE</b>.")
+    )
+
+    with gr.Column():
+        # with gr.Group():
+        gr.Markdown(value=i18n("模型切换"))
+        with gr.Row():
+            GPT_dropdown = gr.Dropdown(label=i18n("GPT模型列表"), choices=sorted(GPT_names, key=custom_sort_key), value=gpt_path, interactive=True)
+            SoVITS_dropdown = gr.Dropdown(label=i18n("SoVITS模型列表"), choices=sorted(SoVITS_names, key=custom_sort_key), value=sovits_path, interactive=True)
+            refresh_button = gr.Button(i18n("刷新模型路径"), variant="primary")
+            refresh_button.click(fn=change_choices, inputs=[], outputs=[SoVITS_dropdown, GPT_dropdown])
+
+
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown(value=i18n("*请上传并填写参考信息"))
+            with gr.Row():
+                inp_ref = gr.Audio(label=i18n("主参考音频(请上传3~10秒内参考音频，超过会报错！)"), type="filepath")
+                inp_refs = gr.File(label=i18n("辅参考音频(可选多个，或不选)"),file_count="multiple")
+            prompt_text = gr.Textbox(label=i18n("主参考音频的文本"), value="", lines=2)
+            with gr.Row():
+                prompt_language = gr.Dropdown(
+                    label=i18n("主参考音频的语种"), choices=list(dict_language.keys()), value=i18n("中文")
+                )
+                with gr.Column():
+                    ref_text_free = gr.Checkbox(label=i18n("开启无参考文本模式。不填参考文本亦相当于开启。"), value=False, interactive=True, show_label=True)
+                    gr.Markdown(i18n("使用无参考文本模式时建议使用微调的GPT，听不清参考音频说的啥(不晓得写啥)可以开，开启后无视填写的参考文本。"))
+
+        with gr.Column():
+            gr.Markdown(value=i18n("*请填写需要合成的目标文本和语种模式"))
+            text = gr.Textbox(label=i18n("需要合成的文本"), value="", lines=20, max_lines=20)
+            text_language = gr.Dropdown(
+                label=i18n("需要合成的文本的语种"), choices=list(dict_language.keys()), value=i18n("中文")
+            )
+
+
+    with gr.Group():
+        gr.Markdown(value=i18n("推理设置"))
+        with gr.Row():
+
+            with gr.Column():
+                batch_size = gr.Slider(minimum=1,maximum=200,step=1,label=i18n("batch_size"),value=20,interactive=True)
+                fragment_interval = gr.Slider(minimum=0.01,maximum=1,step=0.01,label=i18n("分段间隔(秒)"),value=0.3,interactive=True)
+                speed_factor = gr.Slider(minimum=0.6,maximum=1.65,step=0.05,label="speed_factor",value=1.0,interactive=True)
+                top_k = gr.Slider(minimum=1,maximum=100,step=1,label=i18n("top_k"),value=5,interactive=True)
+                top_p = gr.Slider(minimum=0,maximum=1,step=0.05,label=i18n("top_p"),value=1,interactive=True)
+                temperature = gr.Slider(minimum=0,maximum=1,step=0.05,label=i18n("temperature"),value=1,interactive=True)
+                repetition_penalty = gr.Slider(minimum=0,maximum=2,step=0.05,label=i18n("重复惩罚"),value=1.35,interactive=True)
+            with gr.Column():
+                with gr.Row():
+                    how_to_cut = gr.Dropdown(
+                            label=i18n("怎么切"),
+                            choices=[i18n("不切"), i18n("凑四句一切"), i18n("凑50字一切"), i18n("按中文句号。切"), i18n("按英文句号.切"), i18n("按标点符号切"), ],
+                            value=i18n("凑四句一切"),
+                            interactive=True, scale=1
+                        )
+                    parallel_infer = gr.Checkbox(label=i18n("并行推理"), value=True, interactive=True, show_label=True)
+                    split_bucket = gr.Checkbox(label=i18n("数据分桶(并行推理时会降低一点计算量)"), value=True, interactive=True, show_label=True)
+
+                with gr.Row():
+                    seed = gr.Number(label=i18n("随机种子"),value=-1)
+                    keep_random = gr.Checkbox(label=i18n("保持随机"), value=True, interactive=True, show_label=True)
+
+                output = gr.Audio(label=i18n("输出的语音"))
+                with gr.Row():
+                    inference_button = gr.Button(i18n("合成语音"), variant="primary")
+                    stop_infer = gr.Button(i18n("终止合成"), variant="primary")
+
+
+        inference_button.click(
+            inference,
+            [
+                text,text_language, inp_ref, inp_refs,
+                prompt_text, prompt_language,
+                top_k, top_p, temperature,
+                how_to_cut, batch_size,
+                speed_factor, ref_text_free,
+                split_bucket,fragment_interval,
+                seed, keep_random, parallel_infer,
+                repetition_penalty
+             ],
+            [output, seed],
+        )
+        stop_infer.click(tts_pipeline.stop, [], [])
+        SoVITS_dropdown.change(change_sovits_weights, [SoVITS_dropdown,prompt_language,text_language], [prompt_language,text_language,prompt_text,prompt_language,text,text_language])
+        GPT_dropdown.change(tts_pipeline.init_t2s_weights, [GPT_dropdown], [])
+
+    with gr.Group():
+        gr.Markdown(value=i18n("文本切分工具。太长的文本合成出来效果不一定好，所以太长建议先切。合成会根据文本的换行分开合成再拼起来。"))
+        with gr.Row():
+            text_inp = gr.Textbox(label=i18n("需要合成的切分前文本"), value="", lines=4)
+            with gr.Column():
+                _how_to_cut = gr.Radio(
+                            label=i18n("怎么切"),
+                            choices=[i18n("不切"), i18n("凑四句一切"), i18n("凑50字一切"), i18n("按中文句号。切"), i18n("按英文句号.切"), i18n("按标点符号切"), ],
+                            value=i18n("凑四句一切"),
+                            interactive=True,
+                        )
+                cut_text= gr.Button(i18n("切分"), variant="primary")
+
+            def to_cut(text_inp, how_to_cut):
+                if len(text_inp.strip()) == 0 or text_inp==[]:
+                    return ""
+                method = get_method(cut_method[how_to_cut])
+                return method(text_inp)
+
+            text_opt = gr.Textbox(label=i18n("切分后文本"), value="", lines=4)
+            cut_text.click(to_cut, [text_inp, _how_to_cut], [text_opt])
+        gr.Markdown(value=i18n("后续将支持转音素、手工修改音素、语音合成分步执行。"))
+
+if __name__ == '__main__':
+    app.queue().launch(#concurrency_count=511, max_size=1022
+        server_name="0.0.0.0",
+        inbrowser=True,
+        share=is_share,
+        server_port=infer_ttswebui,
+        quiet=True,
+    )

module/attentions.py

@@ -0,0 +1,709 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from moyoyo_tts.module import commons
+from moyoyo_tts.module.modules import LayerNorm
+
+
+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=False,
+        **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
+
+        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))
+        if isflow:
+            cond_layer = torch.nn.Conv1d(
+                kwargs["gin_channels"], 2 * hidden_channels * n_layers, 1
+            )
+            self.cond_pre = torch.nn.Conv1d(hidden_channels, 2 * hidden_channels, 1)
+            self.cond_layer = weight_norm_modules(cond_layer, name="weight")
+            self.gin_channels = kwargs["gin_channels"]
+
+    def forward(self, x, x_mask, g=None):
+        attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        if g is not None:
+            g = self.cond_layer(g)
+
+        for i in range(self.n_layers):
+            if g is not None:
+                x = self.cond_pre(x)
+                cond_offset = i * 2 * self.hidden_channels
+                g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
+                x = commons.fused_add_tanh_sigmoid_multiply(
+                    x, g_l, torch.IntTensor([self.hidden_channels])
+                )
+            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 Decoder(nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size=1,
+        p_dropout=0.0,
+        proximal_bias=False,
+        proximal_init=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.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+
+        self.drop = nn.Dropout(p_dropout)
+        self.self_attn_layers = nn.ModuleList()
+        self.norm_layers_0 = nn.ModuleList()
+        self.encdec_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.self_attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    proximal_bias=proximal_bias,
+                    proximal_init=proximal_init,
+                )
+            )
+            self.norm_layers_0.append(LayerNorm(hidden_channels))
+            self.encdec_attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout
+                )
+            )
+            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,
+                    causal=True,
+                )
+            )
+            self.norm_layers_2.append(LayerNorm(hidden_channels))
+
+    def forward(self, x, x_mask, h, h_mask):
+        """
+        x: decoder input
+        h: encoder output
+        """
+        self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(
+            device=x.device, dtype=x.dtype
+        )
+        encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            y = self.self_attn_layers[i](x, x, self_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_0[i](x + y)
+
+            y = self.encdec_attn_layers[i](x, h, encdec_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, t_t = (*key.size(), query.size(2))
+        query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
+        key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+        value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+
+        scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
+        if self.window_size is not None:
+            assert (
+                t_s == t_t
+            ), "Relative attention is only available for self-attention."
+            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 self.proximal_bias:
+            assert t_s == t_t, "Proximal bias is only available for self-attention."
+            scores = scores + self._attention_bias_proximal(t_s).to(
+                device=scores.device, dtype=scores.dtype
+            )
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e4)
+            if self.block_length is not None:
+                assert (
+                    t_s == t_t
+                ), "Local attention is only available for self-attention."
+                block_mask = (
+                    torch.ones_like(scores)
+                    .triu(-self.block_length)
+                    .tril(self.block_length)
+                )
+                scores = scores.masked_fill(block_mask == 0, -1e4)
+        p_attn = F.softmax(scores, dim=-1)  # [b, n_h, t_t, t_s]
+        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, t_t)
+        )  # [b, n_h, t_t, d_k] -> [b, d, t_t]
+        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_length = max(length - (self.window_size + 1), 0)
+        slice_start_position = max((self.window_size + 1) - length, 0)
+        slice_end_position = slice_start_position + 2 * length - 1
+        if pad_length > 0:
+            padded_relative_embeddings = F.pad(
+                relative_embeddings,
+                commons.convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]),
+            )
+        else:
+            padded_relative_embeddings = relative_embeddings
+        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
+
+
+import torch.nn as nn
+from torch.nn.utils import remove_weight_norm, weight_norm
+
+
+class Depthwise_Separable_Conv1D(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        dilation=1,
+        bias=True,
+        padding_mode="zeros",  # TODO: refine this type
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        self.depth_conv = nn.Conv1d(
+            in_channels=in_channels,
+            out_channels=in_channels,
+            kernel_size=kernel_size,
+            groups=in_channels,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+            padding_mode=padding_mode,
+            device=device,
+            dtype=dtype,
+        )
+        self.point_conv = nn.Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            bias=bias,
+            device=device,
+            dtype=dtype,
+        )
+
+    def forward(self, input):
+        return self.point_conv(self.depth_conv(input))
+
+    def weight_norm(self):
+        self.depth_conv = weight_norm(self.depth_conv, name="weight")
+        self.point_conv = weight_norm(self.point_conv, name="weight")
+
+    def remove_weight_norm(self):
+        self.depth_conv = remove_weight_norm(self.depth_conv, name="weight")
+        self.point_conv = remove_weight_norm(self.point_conv, name="weight")
+
+
+class Depthwise_Separable_TransposeConv1D(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        padding=0,
+        output_padding=0,
+        bias=True,
+        dilation=1,
+        padding_mode="zeros",  # TODO: refine this type
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        self.depth_conv = nn.ConvTranspose1d(
+            in_channels=in_channels,
+            out_channels=in_channels,
+            kernel_size=kernel_size,
+            groups=in_channels,
+            stride=stride,
+            output_padding=output_padding,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+            padding_mode=padding_mode,
+            device=device,
+            dtype=dtype,
+        )
+        self.point_conv = nn.Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            bias=bias,
+            device=device,
+            dtype=dtype,
+        )
+
+    def forward(self, input):
+        return self.point_conv(self.depth_conv(input))
+
+    def weight_norm(self):
+        self.depth_conv = weight_norm(self.depth_conv, name="weight")
+        self.point_conv = weight_norm(self.point_conv, name="weight")
+
+    def remove_weight_norm(self):
+        remove_weight_norm(self.depth_conv, name="weight")
+        remove_weight_norm(self.point_conv, name="weight")
+
+
+def weight_norm_modules(module, name="weight", dim=0):
+    if isinstance(module, Depthwise_Separable_Conv1D) or isinstance(
+        module, Depthwise_Separable_TransposeConv1D
+    ):
+        module.weight_norm()
+        return module
+    else:
+        return weight_norm(module, name, dim)
+
+
+def remove_weight_norm_modules(module, name="weight"):
+    if isinstance(module, Depthwise_Separable_Conv1D) or isinstance(
+        module, Depthwise_Separable_TransposeConv1D
+    ):
+        module.remove_weight_norm()
+    else:
+        remove_weight_norm(module, name)
+
+
+class FFT(nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers=1,
+        kernel_size=1,
+        p_dropout=0.0,
+        proximal_bias=False,
+        proximal_init=True,
+        isflow=False,
+        **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.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+        if isflow:
+            cond_layer = torch.nn.Conv1d(
+                kwargs["gin_channels"], 2 * hidden_channels * n_layers, 1
+            )
+            self.cond_pre = torch.nn.Conv1d(hidden_channels, 2 * hidden_channels, 1)
+            self.cond_layer = weight_norm_modules(cond_layer, name="weight")
+            self.gin_channels = kwargs["gin_channels"]
+        self.drop = nn.Dropout(p_dropout)
+        self.self_attn_layers = nn.ModuleList()
+        self.norm_layers_0 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        for i in range(self.n_layers):
+            self.self_attn_layers.append(
+                MultiHeadAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    proximal_bias=proximal_bias,
+                    proximal_init=proximal_init,
+                )
+            )
+            self.norm_layers_0.append(LayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FFN(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                    causal=True,
+                )
+            )
+            self.norm_layers_1.append(LayerNorm(hidden_channels))
+
+    def forward(self, x, x_mask, g=None):
+        """
+        x: decoder input
+        h: encoder output
+        """
+        if g is not None:
+            g = self.cond_layer(g)
+
+        self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(
+            device=x.device, dtype=x.dtype
+        )
+        x = x * x_mask
+        for i in range(self.n_layers):
+            if g is not None:
+                x = self.cond_pre(x)
+                cond_offset = i * 2 * self.hidden_channels
+                g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
+                x = commons.fused_add_tanh_sigmoid_multiply(
+                    x, g_l, torch.IntTensor([self.hidden_channels])
+                )
+            y = self.self_attn_layers[i](x, x, self_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_0[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+        x = x * x_mask
+        return x
+
+
+class TransformerCouplingLayer(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        n_layers,
+        n_heads,
+        p_dropout=0,
+        filter_channels=0,
+        mean_only=False,
+        wn_sharing_parameter=None,
+        gin_channels=0,
+    ):
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.half_channels = channels // 2
+        self.mean_only = mean_only
+
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = (
+            Encoder(
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers,
+                kernel_size,
+                p_dropout,
+                isflow=True,
+                gin_channels=gin_channels,
+            )
+            if wn_sharing_parameter is None
+            else wn_sharing_parameter
+        )
+        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+        self.post.weight.data.zero_()
+        self.post.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, logs = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            logs = torch.zeros_like(m)
+
+        if not reverse:
+            x1 = m + x1 * torch.exp(logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(logs, [1, 2])
+            return x, logdet
+        else:
+            x1 = (x1 - m) * torch.exp(-logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            return x

module/attentions_onnx.py

@@ -0,0 +1,406 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from moyoyo_tts.module import commons
+
+from typing import Optional
+
+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
+        self.spk_emb_linear = nn.Linear(256, self.hidden_channels)
+        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
+
+    def forward(self, x, x_mask):
+        attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for attn_layers,norm_layers_1,ffn_layers,norm_layers_2 in zip(self.attn_layers,self.norm_layers_1,self.ffn_layers,self.norm_layers_2):
+            y = attn_layers(x, x, attn_mask)
+            y = self.drop(y)
+            x = norm_layers_1(x + y)
+
+            y = ffn_layers(x, x_mask)
+            y = self.drop(y)
+            x = norm_layers_2(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:Optional[torch.Tensor]=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.attention(q, k, v, mask=attn_mask)
+
+        x = self.conv_o(x)
+        return x
+
+    def attention(self, query, key, value, mask:Optional[torch.Tensor]=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="",
+        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
+
+        # 从上下文看这里一定是 False
+        # 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 padding(self, x):
+        return self._same_padding(x)
+
+    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
+
+
+class MRTE(nn.Module):
+    def __init__(
+        self,
+        content_enc_channels=192,
+        hidden_size=512,
+        out_channels=192,
+        kernel_size=5,
+        n_heads=4,
+        ge_layer=2,
+    ):
+        super(MRTE, self).__init__()
+        self.cross_attention = MultiHeadAttention(hidden_size, hidden_size, n_heads)
+        self.c_pre = nn.Conv1d(content_enc_channels, hidden_size, 1)
+        self.text_pre = nn.Conv1d(content_enc_channels, hidden_size, 1)
+        self.c_post = nn.Conv1d(hidden_size, out_channels, 1)
+
+    def forward(self, ssl_enc, ssl_mask, text, text_mask, ge):
+        attn_mask = text_mask.unsqueeze(2) * ssl_mask.unsqueeze(-1)
+
+        ssl_enc = self.c_pre(ssl_enc * ssl_mask)
+        text_enc = self.text_pre(text * text_mask)
+        x = (
+                self.cross_attention(
+                    ssl_enc * ssl_mask, text_enc * text_mask, attn_mask
+                )
+                + ssl_enc
+                + ge
+            )
+        x = self.c_post(x * ssl_mask)
+        return x