初始化项目

.gitignore

@@ -0,0 +1,163 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+*.ckpt
+# C extensions
+*.so
+*.pt
+
+# Distribution / packaging
+.Python
+outputs/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+asset/*
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

ChatTTS/__init__.py

@@ -0,0 +1 @@
+from .core import Chat

ChatTTS/core.py

@@ -0,0 +1,200 @@
+
+import os
+import logging
+from functools import partial
+from omegaconf import OmegaConf
+
+import torch
+from vocos import Vocos
+from .model.dvae import DVAE
+from .model.gpt import GPT_warpper
+from .utils.gpu_utils import select_device
+from .utils.infer_utils import count_invalid_characters, detect_language, apply_character_map, apply_half2full_map
+from .utils.io_utils import get_latest_modified_file
+from .infer.api import refine_text, infer_code
+
+from huggingface_hub import snapshot_download
+
+logging.basicConfig(level = logging.INFO)
+
+
+class Chat:
+    def __init__(self, ):
+        self.pretrain_models = {}
+        self.normalizer = {}
+        self.logger = logging.getLogger(__name__)
+        
+    def check_model(self, level = logging.INFO, use_decoder = False):
+        not_finish = False
+        check_list = ['vocos', 'gpt', 'tokenizer']
+        
+        if use_decoder:
+            check_list.append('decoder')
+        else:
+            check_list.append('dvae')
+            
+        for module in check_list:
+            if module not in self.pretrain_models:
+                self.logger.log(logging.WARNING, f'{module} not initialized.')
+                not_finish = True
+                
+        if not not_finish:
+            self.logger.log(level, f'All initialized.')
+            
+        return not not_finish
+        
+    def load_models(self, source='huggingface', force_redownload=False, local_path='<LOCAL_PATH>', **kwargs):
+        if source == 'huggingface':
+            hf_home = os.getenv('HF_HOME', os.path.expanduser("~/.cache/huggingface"))
+            try:
+                download_path = get_latest_modified_file(os.path.join(hf_home, 'hub/models--2Noise--ChatTTS/snapshots'))
+            except:
+                download_path = None
+            if download_path is None or force_redownload: 
+                self.logger.log(logging.INFO, f'Download from HF: https://huggingface.co/2Noise/ChatTTS')
+                download_path = snapshot_download(repo_id="2Noise/ChatTTS", allow_patterns=["*.pt", "*.yaml"])
+            else:
+                self.logger.log(logging.INFO, f'Load from cache: {download_path}')
+        elif source == 'local':
+            self.logger.log(logging.INFO, f'Load from local: {local_path}')
+            download_path = local_path
+
+        self._load(**{k: os.path.join(download_path, v) for k, v in OmegaConf.load(os.path.join(download_path, 'config', 'path.yaml')).items()}, **kwargs)
+        
+    def _load(
+        self, 
+        vocos_config_path: str = None, 
+        vocos_ckpt_path: str = None,
+        dvae_config_path: str = None,
+        dvae_ckpt_path: str = None,
+        gpt_config_path: str = None,
+        gpt_ckpt_path: str = None,
+        decoder_config_path: str = None,
+        decoder_ckpt_path: str = None,
+        tokenizer_path: str = None,
+        device: str = None,
+        compile: bool = True,
+    ):
+        if not device:
+            device = select_device(4096)
+            self.logger.log(logging.INFO, f'use {device}')
+            
+        if vocos_config_path:
+            vocos = Vocos.from_hparams(vocos_config_path).to(device).eval()
+            assert vocos_ckpt_path, 'vocos_ckpt_path should not be None'
+            vocos.load_state_dict(torch.load(vocos_ckpt_path))
+            self.pretrain_models['vocos'] = vocos
+            self.logger.log(logging.INFO, 'vocos loaded.')
+        
+        if dvae_config_path:
+            cfg = OmegaConf.load(dvae_config_path)
+            dvae = DVAE(**cfg).to(device).eval()
+            assert dvae_ckpt_path, 'dvae_ckpt_path should not be None'
+            dvae.load_state_dict(torch.load(dvae_ckpt_path, map_location='cpu'))
+            self.pretrain_models['dvae'] = dvae
+            self.logger.log(logging.INFO, 'dvae loaded.')
+            
+        if gpt_config_path:
+            cfg = OmegaConf.load(gpt_config_path)
+            gpt = GPT_warpper(**cfg).to(device).eval()
+            assert gpt_ckpt_path, 'gpt_ckpt_path should not be None'
+            gpt.load_state_dict(torch.load(gpt_ckpt_path, map_location='cpu'))
+            if compile and 'cuda' in str(device):
+                gpt.gpt.forward = torch.compile(gpt.gpt.forward,  backend='inductor', dynamic=True)
+            self.pretrain_models['gpt'] = gpt
+            spk_stat_path = os.path.join(os.path.dirname(gpt_ckpt_path), 'spk_stat.pt')
+            assert os.path.exists(spk_stat_path), f'Missing spk_stat.pt: {spk_stat_path}'
+            self.pretrain_models['spk_stat'] = torch.load(spk_stat_path).to(device)
+            self.logger.log(logging.INFO, 'gpt loaded.')
+            
+        if decoder_config_path:
+            cfg = OmegaConf.load(decoder_config_path)
+            decoder = DVAE(**cfg).to(device).eval()
+            assert decoder_ckpt_path, 'decoder_ckpt_path should not be None'
+            decoder.load_state_dict(torch.load(decoder_ckpt_path, map_location='cpu'))
+            self.pretrain_models['decoder'] = decoder
+            self.logger.log(logging.INFO, 'decoder loaded.')
+        
+        if tokenizer_path:
+            tokenizer = torch.load(tokenizer_path, map_location='cpu')
+            tokenizer.padding_side = 'left'
+            self.pretrain_models['tokenizer'] = tokenizer
+            self.logger.log(logging.INFO, 'tokenizer loaded.')
+            
+        self.check_model()
+    
+    def infer(
+        self, 
+        text, 
+        skip_refine_text=False, 
+        refine_text_only=False, 
+        params_refine_text={}, 
+        params_infer_code={'prompt':'[speed_5]'}, 
+        use_decoder=True,
+        do_text_normalization=True,
+        lang=None,
+    ):
+        
+        assert self.check_model(use_decoder=use_decoder)
+        
+        if not isinstance(text, list): 
+            text = [text]
+        
+        # if do_text_normalization:
+        #     for i, t in enumerate(text):
+        #         _lang = detect_language(t) if lang is None else lang
+        #         self.init_normalizer(_lang)
+        #         text[i] = self.normalizer[_lang](t)
+        #         if _lang == 'zh':
+        #             text[i] = apply_half2full_map(text[i])
+            
+        for i, t in enumerate(text):
+            invalid_characters = count_invalid_characters(t)
+            if len(invalid_characters):
+                self.logger.log(logging.WARNING, f'Invalid characters found! : {invalid_characters}')
+                text[i] = apply_character_map(t)
+                
+        if not skip_refine_text:
+            text_tokens = refine_text(self.pretrain_models, text, **params_refine_text)['ids']
+            text_tokens = [i[i < self.pretrain_models['tokenizer'].convert_tokens_to_ids('[break_0]')] for i in text_tokens]
+            text = self.pretrain_models['tokenizer'].batch_decode(text_tokens)
+            if refine_text_only:
+                return text
+            
+        text = [params_infer_code.get('prompt', '') + i for i in text]
+        params_infer_code.pop('prompt', '')
+        result = infer_code(self.pretrain_models, text, **params_infer_code, return_hidden=use_decoder)
+        
+        if use_decoder:
+            mel_spec = [self.pretrain_models['decoder'](i[None].permute(0,2,1)) for i in result['hiddens']]
+        else:
+            mel_spec = [self.pretrain_models['dvae'](i[None].permute(0,2,1)) for i in result['ids']]
+            
+        wav = [self.pretrain_models['vocos'].decode(i).cpu().numpy() for i in mel_spec]
+        
+        return wav
+    
+    def sample_random_speaker(self, ):
+        
+        dim = self.pretrain_models['gpt'].gpt.layers[0].mlp.gate_proj.in_features
+        std, mean = self.pretrain_models['spk_stat'].chunk(2)
+        return torch.randn(dim, device=std.device) * std + mean
+    
+    def init_normalizer(self, lang):
+        
+        if lang not in self.normalizer:
+            if lang == 'zh':
+                try:
+                    from tn.chinese.normalizer import Normalizer
+                except:
+                    self.logger.log(logging.WARNING, f'Package WeTextProcessing not found! \
+                        Run: conda install -c conda-forge pynini=2.1.5 && pip install WeTextProcessing')
+                self.normalizer[lang] = Normalizer().normalize
+            else:
+                try:
+                    from nemo_text_processing.text_normalization.normalize import Normalizer
+                except:
+                    self.logger.log(logging.WARNING, f'Package nemo_text_processing not found! \
+                        Run: conda install -c conda-forge pynini=2.1.5 && pip install nemo_text_processing')
+                self.normalizer[lang] = partial(Normalizer(input_case='cased', lang=lang).normalize, verbose=False, punct_post_process=True)
+

ChatTTS/experimental/llm.py

@@ -0,0 +1,40 @@
+
+from openai import OpenAI
+ 
+prompt_dict = {
+    'kimi': [ {"role": "system", "content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文和英文的对话。"},
+              {"role": "user", "content": "你好，请注意你现在生成的文字要按照人日常生活的口吻，你的回复将会后续用TTS模型转为语音，并且请把回答控制在100字以内。并且标点符号仅包含逗号和句号，将数字等转为文字回答。"},
+              {"role": "assistant", "content": "好的，我现在生成的文字将按照人日常生活的口吻， 并且我会把回答控制在一百字以内, 标点符号仅包含逗号和句号，将阿拉伯数字等转为中文文字回答。下面请开始对话。"},],
+    'deepseek': [
+        {"role": "system", "content": "You are a helpful assistant"},
+        {"role": "user", "content": "你好，请注意你现在生成的文字要按照人日常生活的口吻，你的回复将会后续用TTS模型转为语音，并且请把回答控制在100字以内。并且标点符号仅包含逗号和句号，将数字等转为文字回答。"},
+        {"role": "assistant", "content": "好的，我现在生成的文字将按照人日常生活的口吻， 并且我会把回答控制在一百字以内, 标点符号仅包含逗号和句号，将阿拉伯数字等转为中文文字回答。下面请开始对话。"},],
+    'deepseek_TN': [
+        {"role": "system", "content": "You are a helpful assistant"},
+        {"role": "user", "content": "你好，现在我们在处理TTS的文本输入，下面将会给你输入一段文本，请你将其中的阿拉伯数字等等转为文字表达，并且输出的文本里仅包含逗号和句号这两个标点符号"},
+        {"role": "assistant", "content": "好的，我现在对TTS的文本输入进行处理。这一般叫做text normalization。下面请输入"},
+        {"role": "user", "content": "We paid $123 for this desk."},
+        {"role": "assistant", "content": "We paid one hundred and twenty three dollars for this desk."},
+        {"role": "user", "content": "详询请拨打010-724654"},
+        {"role": "assistant", "content": "详询请拨打零幺零，七二四六五四"},
+        {"role": "user", "content": "罗森宣布将于7月24日退市，在华门店超6000家！"},
+        {"role": "assistant", "content": "罗森宣布将于七月二十四日退市，在华门店超过六千家。"},
+        ],
+}          
+                
+class llm_api:
+    def __init__(self, api_key, base_url, model):
+        self.client =  OpenAI(
+            api_key = api_key,
+            base_url = base_url,
+        )
+        self.model = model
+    def call(self, user_question, temperature = 0.3, prompt_version='kimi', **kwargs):
+    
+        completion = self.client.chat.completions.create(
+            model = self.model,
+            messages = prompt_dict[prompt_version]+[{"role": "user", "content": user_question},],
+            temperature = temperature,
+            **kwargs
+        )
+        return completion.choices[0].message.content

ChatTTS/infer/api.py

@@ -0,0 +1,125 @@
+
+import torch
+import torch.nn.functional as F
+from transformers.generation import TopKLogitsWarper, TopPLogitsWarper
+from ..utils.infer_utils import CustomRepetitionPenaltyLogitsProcessorRepeat
+
+def infer_code(
+    models,
+    text, 
+    spk_emb = None,
+    top_P = 0.7, 
+    top_K = 20, 
+    temperature = 0.3, 
+    repetition_penalty = 1.05,
+    max_new_token = 2048,
+    **kwargs
+):
+    
+    device = next(models['gpt'].parameters()).device
+    
+    if not isinstance(text, list): 
+        text = [text]
+        
+    if not isinstance(temperature, list):
+        temperature = [temperature] * models['gpt'].num_vq
+    
+    if spk_emb is not None:
+        text = [f'[Stts][spk_emb]{i}[Ptts]' for i in text] 
+    else:
+        text = [f'[Stts][empty_spk]{i}[Ptts]' for i in text]
+    
+    text_token = models['tokenizer'](text, return_tensors='pt', add_special_tokens=False, padding=True).to(device)
+    input_ids = text_token['input_ids'][...,None].expand(-1, -1, models['gpt'].num_vq)
+    text_mask = torch.ones(text_token['input_ids'].shape, dtype=bool, device=device)
+    
+    inputs = {
+        'input_ids': input_ids,
+        'text_mask': text_mask,
+        'attention_mask': text_token['attention_mask'],
+    }
+
+    emb = models['gpt'].get_emb(**inputs)
+    if spk_emb is not None:
+        emb[inputs['input_ids'][..., 0] == models['tokenizer'].convert_tokens_to_ids('[spk_emb]')] = \
+            F.normalize(spk_emb.to(device).to(emb.dtype)[None].expand(len(text), -1), p=2.0, dim=1, eps=1e-12)  
+    
+    num_code = models['gpt'].emb_code[0].num_embeddings - 1
+    
+    LogitsWarpers = []
+    if top_P is not None:
+        LogitsWarpers.append(TopPLogitsWarper(top_P, min_tokens_to_keep=3))
+    if top_K is not None:
+        LogitsWarpers.append(TopKLogitsWarper(top_K, min_tokens_to_keep=3))
+        
+    LogitsProcessors = []
+    if repetition_penalty is not None and repetition_penalty != 1:
+        LogitsProcessors.append(CustomRepetitionPenaltyLogitsProcessorRepeat(\
+            repetition_penalty, num_code, 16))
+    
+    result = models['gpt'].generate(
+        emb, inputs['input_ids'], 
+        temperature = torch.tensor(temperature, device=device), 
+        attention_mask = inputs['attention_mask'],
+        LogitsWarpers = LogitsWarpers,
+        LogitsProcessors = LogitsProcessors,
+        eos_token = num_code, 
+        max_new_token = max_new_token, 
+        infer_text = False,
+        **kwargs
+    )
+    
+    return result
+
+
+def refine_text(
+    models, 
+    text,
+    top_P = 0.7, 
+    top_K = 20, 
+    temperature = 0.7, 
+    repetition_penalty = 1.0,
+    max_new_token = 384,
+    prompt = '',
+    **kwargs
+):
+    
+    device = next(models['gpt'].parameters()).device
+    
+    if not isinstance(text, list): 
+        text = [text]
+    
+    assert len(text), 'text should not be empty'
+
+    text = [f"[Sbreak]{i}[Pbreak]{prompt}" for i in text]
+    text_token = models['tokenizer'](text, return_tensors='pt', add_special_tokens=False, padding=True).to(device)
+    text_mask = torch.ones(text_token['input_ids'].shape, dtype=bool, device=device)
+
+    inputs = {
+        'input_ids': text_token['input_ids'][...,None].expand(-1, -1, models['gpt'].num_vq),
+        'text_mask': text_mask,
+        'attention_mask': text_token['attention_mask'],
+    }
+    
+    LogitsWarpers = []
+    if top_P is not None:
+        LogitsWarpers.append(TopPLogitsWarper(top_P, min_tokens_to_keep=3))
+    if top_K is not None:
+        LogitsWarpers.append(TopKLogitsWarper(top_K, min_tokens_to_keep=3))
+        
+    LogitsProcessors = []
+    if repetition_penalty is not None and repetition_penalty != 1:
+        LogitsProcessors.append(CustomRepetitionPenaltyLogitsProcessorRepeat(repetition_penalty, len(models['tokenizer']), 16))
+    
+    result = models['gpt'].generate(
+        models['gpt'].get_emb(**inputs), inputs['input_ids'], 
+        temperature = torch.tensor([temperature,], device=device), 
+        attention_mask = inputs['attention_mask'],
+        LogitsWarpers = LogitsWarpers,
+        LogitsProcessors = LogitsProcessors,
+        eos_token = torch.tensor(models['tokenizer'].convert_tokens_to_ids('[Ebreak]'), device=device)[None], 
+        max_new_token = max_new_token, 
+        infer_text = True,
+        **kwargs
+    )
+    return result

ChatTTS/model/dvae.py

@@ -0,0 +1,155 @@
+import math
+from einops import rearrange
+from vector_quantize_pytorch import GroupedResidualFSQ
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class ConvNeXtBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        intermediate_dim: int,
+        kernel, dilation,
+        layer_scale_init_value: float = 1e-6,
+    ):
+        # ConvNeXt Block copied from Vocos.
+        super().__init__()
+        self.dwconv = nn.Conv1d(dim, dim, 
+                                kernel_size=kernel, padding=dilation*(kernel//2), 
+                                dilation=dilation, groups=dim
+                            )  # depthwise conv
+        
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, intermediate_dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(intermediate_dim, dim)
+        self.gamma = (
+            nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
+            if layer_scale_init_value > 0
+            else None
+        )
+
+    def forward(self, x: torch.Tensor, cond = None) -> torch.Tensor:
+        residual = x
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)  # (B, C, T) -> (B, T, C)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.transpose(1, 2)  # (B, T, C) -> (B, C, T)
+
+        x = residual + x
+        return x
+    
+
+
+class GFSQ(nn.Module):
+
+    def __init__(self, 
+            dim, levels, G, R, eps=1e-5, transpose = True
+        ):
+        super(GFSQ, self).__init__()
+        self.quantizer = GroupedResidualFSQ(
+            dim=dim,
+            levels=levels,
+            num_quantizers=R,
+            groups=G,
+        )
+        self.n_ind = math.prod(levels)
+        self.eps = eps
+        self.transpose = transpose
+        self.G = G
+        self.R = R
+        
+    def _embed(self, x):
+        if self.transpose:
+            x = x.transpose(1,2)
+        x = rearrange(
+            x, "b t (g r) -> g b t r", g = self.G, r = self.R,
+        )  
+        feat = self.quantizer.get_output_from_indices(x)
+        return feat.transpose(1,2) if self.transpose else feat
+        
+    def forward(self, x,):
+        if self.transpose:
+            x = x.transpose(1,2)
+        feat, ind = self.quantizer(x)
+        ind = rearrange(
+            ind, "g b t r ->b t (g r)",
+        )  
+        embed_onehot = F.one_hot(ind.long(), self.n_ind).to(x.dtype)
+        e_mean = torch.mean(embed_onehot, dim=[0,1])
+        e_mean = e_mean / (e_mean.sum(dim=1) + self.eps).unsqueeze(1)
+        perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + self.eps), dim=1))
+        
+        return (
+            torch.zeros(perplexity.shape, dtype=x.dtype, device=x.device),
+            feat.transpose(1,2) if self.transpose else feat,
+            perplexity,
+            None,
+            ind.transpose(1,2) if self.transpose else ind,
+        )
+        
+class DVAEDecoder(nn.Module):
+    def __init__(self, idim, odim,
+                 n_layer = 12, bn_dim = 64, hidden = 256, 
+                 kernel = 7, dilation = 2, up = False
+                ):
+        super().__init__()
+        self.up = up
+        self.conv_in = nn.Sequential(
+            nn.Conv1d(idim, bn_dim, 3, 1, 1), nn.GELU(),
+            nn.Conv1d(bn_dim, hidden, 3, 1, 1)
+        )
+        self.decoder_block = nn.ModuleList([
+            ConvNeXtBlock(hidden, hidden* 4, kernel, dilation,)
+            for _ in range(n_layer)])
+        self.conv_out = nn.Conv1d(hidden, odim, kernel_size=1, bias=False)
+
+    def forward(self, input, conditioning=None):
+        # B, T, C
+        x = input.transpose(1, 2)
+        x = self.conv_in(x)
+        for f in self.decoder_block:
+            x = f(x, conditioning)
+        
+        x = self.conv_out(x)
+        return x.transpose(1, 2)
+    
+
+class DVAE(nn.Module):
+    def __init__(
+        self, decoder_config, vq_config, dim=512
+    ):
+        super().__init__()
+        self.register_buffer('coef', torch.randn(1, 100, 1))
+
+        self.decoder = DVAEDecoder(**decoder_config)
+        self.out_conv = nn.Conv1d(dim, 100, 3, 1, 1, bias=False)
+        if vq_config is not None:
+            self.vq_layer = GFSQ(**vq_config)
+        else:
+            self.vq_layer = None
+
+    def forward(self, inp):
+
+        if self.vq_layer is not None:
+            vq_feats = self.vq_layer._embed(inp)
+        else:
+            vq_feats = inp.detach().clone()
+            
+        temp = torch.chunk(vq_feats, 2, dim=1) # flatten trick :)
+        temp = torch.stack(temp, -1)
+        vq_feats = temp.reshape(*temp.shape[:2], -1)
+
+        vq_feats = vq_feats.transpose(1, 2)
+        dec_out = self.decoder(input=vq_feats)
+        dec_out = self.out_conv(dec_out.transpose(1, 2))
+        mel = dec_out * self.coef
+
+        return mel

ChatTTS/model/gpt.py

@@ -0,0 +1,265 @@
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+import logging
+from tqdm import tqdm
+from einops import rearrange
+from transformers.cache_utils import Cache
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.utils.parametrize as P
+from torch.nn.utils.parametrizations import weight_norm
+from transformers import LlamaModel, LlamaConfig
+    
+    
+class LlamaMLP(nn.Module):
+    def __init__(self, hidden_size, intermediate_size):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = F.silu
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+    
+    
+class GPT_warpper(nn.Module):
+    def __init__(
+        self, 
+        gpt_config, 
+        num_audio_tokens,
+        num_text_tokens,
+        num_vq=4,
+        **kwargs,
+        ):
+        super().__init__()
+
+        self.logger = logging.getLogger(__name__)
+        self.gpt = self.build_model(gpt_config)
+        self.model_dim = self.gpt.config.hidden_size 
+
+        self.num_vq = num_vq
+        self.emb_code = nn.ModuleList([nn.Embedding(num_audio_tokens, self.model_dim) for i in range(self.num_vq)])
+        self.emb_text = nn.Embedding(num_text_tokens, self.model_dim)
+        self.head_text = weight_norm(nn.Linear(self.model_dim, num_text_tokens, bias=False), name='weight')
+        self.head_code = nn.ModuleList([weight_norm(nn.Linear(self.model_dim, num_audio_tokens, bias=False), name='weight') for i in range(self.num_vq)])
+
+    def build_model(self, config):
+        
+        configuration = LlamaConfig(**config)
+        model = LlamaModel(configuration)
+        del model.embed_tokens
+        
+        return model
+    
+    def get_emb(self, input_ids, text_mask, **kwargs):
+
+        emb_text = self.emb_text(input_ids[text_mask][:, 0])
+        
+        emb_code = [self.emb_code[i](input_ids[~text_mask][:, i]) for i in range(self.num_vq)]
+        emb_code = torch.stack(emb_code, 2).sum(2)
+        
+        emb = torch.zeros((input_ids.shape[:-1])+(emb_text.shape[-1],), device=emb_text.device, dtype=emb_text.dtype)
+        emb[text_mask] = emb_text
+        emb[~text_mask] = emb_code.to(emb.dtype)
+        
+        return emb
+    
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, cache_position=None, **kwargs
+    ):
+        # With static cache, the `past_key_values` is None
+        # TODO joao: standardize interface for the different Cache classes and remove of this if
+        has_static_cache = False
+        if past_key_values is None:
+            past_key_values = getattr(self.gpt.layers[0].self_attn, "past_key_value", None)
+            has_static_cache = past_key_values is not None
+
+        past_length = 0
+        if past_key_values is not None:
+            if isinstance(past_key_values, Cache):
+                past_length = cache_position[0] if cache_position is not None else past_key_values.get_seq_length()
+                max_cache_length = (
+                    torch.tensor(past_key_values.get_max_length(), device=input_ids.device)
+                    if past_key_values.get_max_length() is not None
+                    else None
+                )
+                cache_length = past_length if max_cache_length is None else torch.min(max_cache_length, past_length)
+            # TODO joao: remove this `else` after `generate` prioritizes `Cache` objects
+            else:
+                cache_length = past_length = past_key_values[0][0].shape[2]
+                max_cache_length = None
+
+            # Keep only the unprocessed tokens:
+            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
+            # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
+            # input)
+            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
+                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
+            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
+            # input_ids based on the past_length.
+            elif past_length < input_ids.shape[1]:
+                input_ids = input_ids[:, past_length:]
+            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
+
+            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
+            if (
+                max_cache_length is not None
+                and attention_mask is not None
+                and cache_length + input_ids.shape[1] > max_cache_length
+            ):
+                attention_mask = attention_mask[:, -max_cache_length:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -input_ids.shape[1] :]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
+            # recompiles graphs as the stride of the inputs is a guard. Ref: https://github.com/huggingface/transformers/pull/29114
+            # TODO: use `next_tokens` directly instead.
+            model_inputs = {"input_ids": input_ids.contiguous()}
+
+        input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1]
+        if cache_position is None:
+            cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device)
+        else:
+            cache_position = cache_position[-input_length:]
+
+        if has_static_cache:
+            past_key_values = None
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+    
+    def generate(
+        self, 
+        emb, 
+        inputs_ids, 
+        temperature, 
+        eos_token, 
+        attention_mask = None,
+        max_new_token = 2048, 
+        min_new_token = 0,
+        LogitsWarpers = [],
+        LogitsProcessors = [],
+        infer_text=False,
+        return_attn=False,
+        return_hidden=False,
+    ):
+        
+        with torch.no_grad():   
+        
+            attentions = []
+            hiddens = []
+            
+            start_idx, end_idx = inputs_ids.shape[1], torch.zeros(inputs_ids.shape[0], device=inputs_ids.device, dtype=torch.long)
+            finish = torch.zeros(inputs_ids.shape[0], device=inputs_ids.device).bool()
+            
+            temperature = temperature[None].expand(inputs_ids.shape[0], -1)
+            temperature = rearrange(temperature, "b n -> (b n) 1")
+
+            attention_mask_cache = torch.ones((inputs_ids.shape[0], inputs_ids.shape[1]+max_new_token,), dtype=torch.bool, device=inputs_ids.device)
+            if attention_mask is not None:
+                attention_mask_cache[:, :attention_mask.shape[1]] = attention_mask
+            
+            for i in tqdm(range(max_new_token)):
+        
+                model_input = self.prepare_inputs_for_generation(inputs_ids, 
+                    outputs.past_key_values if i!=0 else None, 
+                    attention_mask_cache[:, :inputs_ids.shape[1]], use_cache=True)
+            
+                if i == 0:
+                    model_input['inputs_embeds'] = emb
+                else:
+                    if infer_text:
+                        model_input['inputs_embeds'] = self.emb_text(model_input['input_ids'][:,:,0])
+                    else:
+                        code_emb = [self.emb_code[i](model_input['input_ids'][:,:,i]) for i in range(self.num_vq)]
+                        model_input['inputs_embeds'] = torch.stack(code_emb, 3).sum(3)
+                
+                model_input['input_ids'] = None
+                outputs = self.gpt.forward(**model_input, output_attentions=return_attn)
+                attentions.append(outputs.attentions)
+                hidden_states = outputs[0] # 🐻
+                if return_hidden:
+                    hiddens.append(hidden_states[:, -1])
+
+                with P.cached():
+                    if infer_text:
+                        logits = self.head_text(hidden_states) 
+                    else:
+                        logits = torch.stack([self.head_code[i](hidden_states) for i in range(self.num_vq)], 3)
+        
+                logits = logits[:, -1].float()
+
+                if not infer_text:
+                    logits = rearrange(logits, "b c n -> (b n) c")
+                    logits_token = rearrange(inputs_ids[:, start_idx:], "b c n -> (b n) c")
+                else:
+                    logits_token = inputs_ids[:, start_idx:, 0]
+                    
+                logits = logits / temperature
+                
+                for logitsProcessors in LogitsProcessors:
+                    logits = logitsProcessors(logits_token, logits)
+                    
+                for logitsWarpers in LogitsWarpers:
+                    logits = logitsWarpers(logits_token, logits)
+                    
+                if i < min_new_token:
+                    logits[:, eos_token] = -torch.inf
+                
+                scores = F.softmax(logits, dim=-1)
+            
+                idx_next = torch.multinomial(scores, num_samples=1)
+                
+                if not infer_text:
+                    idx_next = rearrange(idx_next, "(b n) 1 -> b n", n=self.num_vq)
+                    finish = finish | (idx_next == eos_token).any(1)
+                    inputs_ids = torch.cat([inputs_ids, idx_next.unsqueeze(1)], 1)
+                else:
+                    finish = finish | (idx_next == eos_token).any(1)
+                    inputs_ids = torch.cat([inputs_ids, idx_next.unsqueeze(-1).expand(-1, -1, self.num_vq)], 1)
+
+                end_idx = end_idx + (~finish).int()
+            
+                if finish.all():
+                    break
+            
+            inputs_ids = [inputs_ids[idx, start_idx: start_idx+i] for idx, i in enumerate(end_idx.int())]
+            inputs_ids = [i[:, 0] for i in inputs_ids] if infer_text else inputs_ids
+            
+            if return_hidden:
+                hiddens = torch.stack(hiddens, 1)
+                hiddens = [hiddens[idx, :i] for idx, i in enumerate(end_idx.int())]
+                    
+            if not finish.all():
+                self.logger.warn(f'Incomplete result. hit max_new_token: {max_new_token}')    
+                   
+            return {
+                'ids': inputs_ids, 
+                'attentions': attentions,
+                'hiddens':hiddens,
+            }

ChatTTS/utils/gpu_utils.py

@@ -0,0 +1,23 @@
+
+import torch
+import logging
+
+def select_device(min_memory = 2048):
+    logger = logging.getLogger(__name__)
+    if torch.cuda.is_available():
+        available_gpus = []
+        for i in range(torch.cuda.device_count()):
+            props = torch.cuda.get_device_properties(i)
+            free_memory = props.total_memory - torch.cuda.memory_reserved(i)
+            available_gpus.append((i, free_memory))
+        selected_gpu, max_free_memory = max(available_gpus, key=lambda x: x[1])
+        device = torch.device(f'cuda:{selected_gpu}')
+        free_memory_mb = max_free_memory / (1024 * 1024)
+        if free_memory_mb < min_memory:
+            logger.log(logging.WARNING, f'GPU {selected_gpu} has {round(free_memory_mb, 2)} MB memory left.')
+            device = torch.device('cpu')
+    else:
+        logger.log(logging.WARNING, f'No GPU found, use CPU instead')
+        device = torch.device('cpu')
+    
+    return device

ChatTTS/utils/infer_utils.py

@@ -0,0 +1,141 @@
+
+import re
+import torch
+import torch.nn.functional as F
+
+    
+class CustomRepetitionPenaltyLogitsProcessorRepeat():
+
+    def __init__(self, penalty: float, max_input_ids, past_window):
+        if not isinstance(penalty, float) or not (penalty > 0):
+            raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}")
+
+        self.penalty = penalty
+        self.max_input_ids = max_input_ids
+        self.past_window = past_window
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        
+        input_ids = input_ids[:, -self.past_window:]
+        freq = F.one_hot(input_ids, scores.size(1)).sum(1)
+        freq[self.max_input_ids:] = 0
+        alpha = self.penalty**freq
+        scores = torch.where(scores < 0, scores*alpha, scores/alpha)
+
+        return scores
+    
+class CustomRepetitionPenaltyLogitsProcessor():
+
+    def __init__(self, penalty: float, max_input_ids, past_window):
+        if not isinstance(penalty, float) or not (penalty > 0):
+            raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}")
+
+        self.penalty = penalty
+        self.max_input_ids = max_input_ids
+        self.past_window = past_window
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        
+        input_ids = input_ids[:, -self.past_window:]
+        score = torch.gather(scores, 1, input_ids)
+        _score = score.detach().clone()
+        score = torch.where(score < 0, score * self.penalty, score / self.penalty)
+        score[input_ids>=self.max_input_ids] = _score[input_ids>=self.max_input_ids]
+        scores.scatter_(1, input_ids, score)
+        
+        return scores
+    
+def count_invalid_characters(s):
+    
+    s = re.sub(r'\[uv_break\]|\[laugh\]|\[lbreak\]', '', s)
+    pattern = re.compile(r'[^\u4e00-\u9fffA-Za-z，。、,\. ]')
+    non_alphabetic_chinese_chars = pattern.findall(s)
+    return set(non_alphabetic_chinese_chars)
+
+def detect_language(sentence):
+
+    chinese_char_pattern = re.compile(r'[\u4e00-\u9fff]')
+    english_word_pattern = re.compile(r'\b[A-Za-z]+\b')
+
+    chinese_chars = chinese_char_pattern.findall(sentence)
+    english_words = english_word_pattern.findall(sentence)
+
+    if len(chinese_chars) > len(english_words):
+        return "zh"
+    else:
+        return "en"
+    
+    
+character_map = {
+    '：': '，',
+    '；': '，',
+    '！': '。',
+    '（': '，',
+    '）': '，',
+    '【': '，',
+    '】': '，',
+    '『': '，',
+    '』': '，',
+    '「': '，',
+    '」': '，',
+    '《': '，',
+    '》': '，',
+    '－': '，',
+    '‘': '',
+    '“': '',
+    '’': '',
+    '”': '',
+    ':': ',',
+    ';': ',',
+    '!': '.',
+    '(': ',',
+    ')': ',',
+    '[': ',',
+    ']': ',',
+    '>': ',',
+    '<': ',',
+    '-': ',',
+}
+
+halfwidth_2_fullwidth_map = {
+        '!': '！',
+        '"': '“',
+        "'": '‘',
+        '#': '＃',
+        '$': '＄',
+        '%': '％',
+        '&': '＆',
+        '(': '（',
+        ')': '）',
+        ',': '，',
+        '-': '－',
+        '*': '＊',
+        '+': '＋',
+        '.': '。',
+        '/': '／',
+        ':': '：',
+        ';': '；',
+        '<': '＜',
+        '=': '＝',
+        '>': '＞',
+        '?': '？',
+        '@': '＠',
+        # '[': '［',
+        '\\': '＼',
+        # ']': '］',
+        '^': '＾',
+        # '_': '＿',
+        '`': '｀',
+        '{': '｛',
+        '|': '｜',
+        '}': '｝',
+        '~': '～'
+    }
+
+def apply_half2full_map(text):
+    translation_table = str.maketrans(halfwidth_2_fullwidth_map)
+    return text.translate(translation_table)
+
+def apply_character_map(text):
+    translation_table = str.maketrans(character_map)
+    return text.translate(translation_table)

ChatTTS/utils/io_utils.py

@@ -0,0 +1,14 @@
+
+import os
+import logging
+
+def get_latest_modified_file(directory):
+    logger = logging.getLogger(__name__)
+    
+    files = [os.path.join(directory, f) for f in os.listdir(directory)] 
+    if not files:
+        logger.log(logging.WARNING, f'No files found in the directory: {directory}')
+        return None
+    latest_file = max(files, key=os.path.getmtime)
+
+    return latest_file

LICENSE

@@ -0,0 +1,407 @@
+Attribution-NonCommercial 4.0 International
+
+=======================================================================
+
+Creative Commons Corporation ("Creative Commons") is not a law firm and
+does not provide legal services or legal advice. Distribution of
+Creative Commons public licenses does not create a lawyer-client or
+other relationship. Creative Commons makes its licenses and related
+information available on an "as-is" basis. Creative Commons gives no
+warranties regarding its licenses, any material licensed under their
+terms and conditions, or any related information. Creative Commons
+disclaims all liability for damages resulting from their use to the
+fullest extent possible.
+
+Using Creative Commons Public Licenses
+
+Creative Commons public licenses provide a standard set of terms and
+conditions that creators and other rights holders may use to share
+original works of authorship and other material subject to copyright
+and certain other rights specified in the public license below. The
+following considerations are for informational purposes only, are not
+exhaustive, and do not form part of our licenses.
+
+     Considerations for licensors: Our public licenses are
+     intended for use by those authorized to give the public
+     permission to use material in ways otherwise restricted by
+     copyright and certain other rights. Our licenses are
+     irrevocable. Licensors should read and understand the terms
+     and conditions of the license they choose before applying it.
+     Licensors should also secure all rights necessary before
+     applying our licenses so that the public can reuse the
+     material as expected. Licensors should clearly mark any
+     material not subject to the license. This includes other CC-
+     licensed material, or material used under an exception or
+     limitation to copyright. More considerations for licensors:
+    wiki.creativecommons.org/Considerations_for_licensors
+
+     Considerations for the public: By using one of our public
+     licenses, a licensor grants the public permission to use the
+     licensed material under specified terms and conditions. If
+     the licensor's permission is not necessary for any reason--for
+     example, because of any applicable exception or limitation to
+     copyright--then that use is not regulated by the license. Our
+     licenses grant only permissions under copyright and certain
+     other rights that a licensor has authority to grant. Use of
+     the licensed material may still be restricted for other
+     reasons, including because others have copyright or other
+     rights in the material. A licensor may make special requests,
+     such as asking that all changes be marked or described.
+     Although not required by our licenses, you are encouraged to
+     respect those requests where reasonable. More considerations
+     for the public:
+    wiki.creativecommons.org/Considerations_for_licensees
+
+=======================================================================
+
+Creative Commons Attribution-NonCommercial 4.0 International Public
+License
+
+By exercising the Licensed Rights (defined below), You accept and agree
+to be bound by the terms and conditions of this Creative Commons
+Attribution-NonCommercial 4.0 International Public License ("Public
+License"). To the extent this Public License may be interpreted as a
+contract, You are granted the Licensed Rights in consideration of Your
+acceptance of these terms and conditions, and the Licensor grants You
+such rights in consideration of benefits the Licensor receives from
+making the Licensed Material available under these terms and
+conditions.
+
+
+Section 1 -- Definitions.
+
+  a. Adapted Material means material subject to Copyright and Similar
+     Rights that is derived from or based upon the Licensed Material
+     and in which the Licensed Material is translated, altered,
+     arranged, transformed, or otherwise modified in a manner requiring
+     permission under the Copyright and Similar Rights held by the
+     Licensor. For purposes of this Public License, where the Licensed
+     Material is a musical work, performance, or sound recording,
+     Adapted Material is always produced where the Licensed Material is
+     synched in timed relation with a moving image.
+
+  b. Adapter's License means the license You apply to Your Copyright
+     and Similar Rights in Your contributions to Adapted Material in
+     accordance with the terms and conditions of this Public License.
+
+  c. Copyright and Similar Rights means copyright and/or similar rights
+     closely related to copyright including, without limitation,
+     performance, broadcast, sound recording, and Sui Generis Database
+     Rights, without regard to how the rights are labeled or
+     categorized. For purposes of this Public License, the rights
+     specified in Section 2(b)(1)-(2) are not Copyright and Similar
+     Rights.
+  d. Effective Technological Measures means those measures that, in the
+     absence of proper authority, may not be circumvented under laws
+     fulfilling obligations under Article 11 of the WIPO Copyright
+     Treaty adopted on December 20, 1996, and/or similar international
+     agreements.
+
+  e. Exceptions and Limitations means fair use, fair dealing, and/or
+     any other exception or limitation to Copyright and Similar Rights
+     that applies to Your use of the Licensed Material.
+
+  f. Licensed Material means the artistic or literary work, database,
+     or other material to which the Licensor applied this Public
+     License.
+
+  g. Licensed Rights means the rights granted to You subject to the
+     terms and conditions of this Public License, which are limited to
+     all Copyright and Similar Rights that apply to Your use of the
+     Licensed Material and that the Licensor has authority to license.
+
+  h. Licensor means the individual(s) or entity(ies) granting rights
+     under this Public License.
+
+  i. NonCommercial means not primarily intended for or directed towards
+     commercial advantage or monetary compensation. For purposes of
+     this Public License, the exchange of the Licensed Material for
+     other material subject to Copyright and Similar Rights by digital
+     file-sharing or similar means is NonCommercial provided there is
+     no payment of monetary compensation in connection with the
+     exchange.
+
+  j. Share means to provide material to the public by any means or
+     process that requires permission under the Licensed Rights, such
+     as reproduction, public display, public performance, distribution,
+     dissemination, communication, or importation, and to make material
+     available to the public including in ways that members of the
+     public may access the material from a place and at a time
+     individually chosen by them.
+
+  k. Sui Generis Database Rights means rights other than copyright
+     resulting from Directive 96/9/EC of the European Parliament and of
+     the Council of 11 March 1996 on the legal protection of databases,
+     as amended and/or succeeded, as well as other essentially
+     equivalent rights anywhere in the world.
+
+  l. You means the individual or entity exercising the Licensed Rights
+     under this Public License. Your has a corresponding meaning.
+
+
+Section 2 -- Scope.
+
+  a. License grant.
+
+       1. Subject to the terms and conditions of this Public License,
+          the Licensor hereby grants You a worldwide, royalty-free,
+          non-sublicensable, non-exclusive, irrevocable license to
+          exercise the Licensed Rights in the Licensed Material to:
+
+            a. reproduce and Share the Licensed Material, in whole or
+               in part, for NonCommercial purposes only; and
+
+            b. produce, reproduce, and Share Adapted Material for
+               NonCommercial purposes only.
+
+       2. Exceptions and Limitations. For the avoidance of doubt, where
+          Exceptions and Limitations apply to Your use, this Public
+          License does not apply, and You do not need to comply with
+          its terms and conditions.
+
+       3. Term. The term of this Public License is specified in Section
+          6(a).
+
+       4. Media and formats; technical modifications allowed. The
+          Licensor authorizes You to exercise the Licensed Rights in
+          all media and formats whether now known or hereafter created,
+          and to make technical modifications necessary to do so. The
+          Licensor waives and/or agrees not to assert any right or
+          authority to forbid You from making technical modifications
+          necessary to exercise the Licensed Rights, including
+          technical modifications necessary to circumvent Effective
+          Technological Measures. For purposes of this Public License,
+          simply making modifications authorized by this Section 2(a)
+          (4) never produces Adapted Material.
+
+       5. Downstream recipients.
+
+            a. Offer from the Licensor -- Licensed Material. Every
+               recipient of the Licensed Material automatically
+               receives an offer from the Licensor to exercise the
+               Licensed Rights under the terms and conditions of this
+               Public License.
+
+            b. No downstream restrictions. You may not offer or impose
+               any additional or different terms or conditions on, or
+               apply any Effective Technological Measures to, the
+               Licensed Material if doing so restricts exercise of the
+               Licensed Rights by any recipient of the Licensed
+               Material.
+
+       6. No endorsement. Nothing in this Public License constitutes or
+          may be construed as permission to assert or imply that You
+          are, or that Your use of the Licensed Material is, connected
+          with, or sponsored, endorsed, or granted official status by,
+          the Licensor or others designated to receive attribution as
+          provided in Section 3(a)(1)(A)(i).
+
+  b. Other rights.
+
+       1. Moral rights, such as the right of integrity, are not
+          licensed under this Public License, nor are publicity,
+          privacy, and/or other similar personality rights; however, to
+          the extent possible, the Licensor waives and/or agrees not to
+          assert any such rights held by the Licensor to the limited
+          extent necessary to allow You to exercise the Licensed
+          Rights, but not otherwise.
+
+       2. Patent and trademark rights are not licensed under this
+          Public License.
+
+       3. To the extent possible, the Licensor waives any right to
+          collect royalties from You for the exercise of the Licensed
+          Rights, whether directly or through a collecting society
+          under any voluntary or waivable statutory or compulsory
+          licensing scheme. In all other cases the Licensor expressly
+          reserves any right to collect such royalties, including when
+          the Licensed Material is used other than for NonCommercial
+          purposes.
+
+
+Section 3 -- License Conditions.
+
+Your exercise of the Licensed Rights is expressly made subject to the
+following conditions.
+
+  a. Attribution.
+
+       1. If You Share the Licensed Material (including in modified
+          form), You must:
+
+            a. retain the following if it is supplied by the Licensor
+               with the Licensed Material:
+
+                 i. identification of the creator(s) of the Licensed
+                    Material and any others designated to receive
+                    attribution, in any reasonable manner requested by
+                    the Licensor (including by pseudonym if
+                    designated);
+
+                ii. a copyright notice;
+
+               iii. a notice that refers to this Public License;
+
+                iv. a notice that refers to the disclaimer of
+                    warranties;
+
+                 v. a URI or hyperlink to the Licensed Material to the
+                    extent reasonably practicable;
+
+            b. indicate if You modified the Licensed Material and
+               retain an indication of any previous modifications; and
+
+            c. indicate the Licensed Material is licensed under this
+               Public License, and include the text of, or the URI or
+               hyperlink to, this Public License.
+
+       2. You may satisfy the conditions in Section 3(a)(1) in any
+          reasonable manner based on the medium, means, and context in
+          which You Share the Licensed Material. For example, it may be
+          reasonable to satisfy the conditions by providing a URI or
+          hyperlink to a resource that includes the required
+          information.
+
+       3. If requested by the Licensor, You must remove any of the
+          information required by Section 3(a)(1)(A) to the extent
+          reasonably practicable.
+
+       4. If You Share Adapted Material You produce, the Adapter's
+          License You apply must not prevent recipients of the Adapted
+          Material from complying with this Public License.
+
+
+Section 4 -- Sui Generis Database Rights.
+
+Where the Licensed Rights include Sui Generis Database Rights that
+apply to Your use of the Licensed Material:
+
+  a. for the avoidance of doubt, Section 2(a)(1) grants You the right
+     to extract, reuse, reproduce, and Share all or a substantial
+     portion of the contents of the database for NonCommercial purposes
+     only;
+
+  b. if You include all or a substantial portion of the database
+     contents in a database in which You have Sui Generis Database
+     Rights, then the database in which You have Sui Generis Database
+     Rights (but not its individual contents) is Adapted Material; and
+
+  c. You must comply with the conditions in Section 3(a) if You Share
+     all or a substantial portion of the contents of the database.
+
+For the avoidance of doubt, this Section 4 supplements and does not
+replace Your obligations under this Public License where the Licensed
+Rights include other Copyright and Similar Rights.
+
+
+Section 5 -- Disclaimer of Warranties and Limitation of Liability.
+
+  a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
+     EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
+     AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
+     ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
+     IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
+     WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
+     PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
+     ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
+     KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
+     ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
+
+  b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
+     TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
+     NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
+     INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
+     COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
+     USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
+     ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
+     DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
+     IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
+
+  c. The disclaimer of warranties and limitation of liability provided
+     above shall be interpreted in a manner that, to the extent
+     possible, most closely approximates an absolute disclaimer and
+     waiver of all liability.
+
+
+Section 6 -- Term and Termination.
+
+  a. This Public License applies for the term of the Copyright and
+     Similar Rights licensed here. However, if You fail to comply with
+     this Public License, then Your rights under this Public License
+     terminate automatically.
+
+  b. Where Your right to use the Licensed Material has terminated under
+     Section 6(a), it reinstates:
+
+       1. automatically as of the date the violation is cured, provided
+          it is cured within 30 days of Your discovery of the
+          violation; or
+
+       2. upon express reinstatement by the Licensor.
+
+     For the avoidance of doubt, this Section 6(b) does not affect any
+     right the Licensor may have to seek remedies for Your violations
+     of this Public License.
+
+  c. For the avoidance of doubt, the Licensor may also offer the
+     Licensed Material under separate terms or conditions or stop
+     distributing the Licensed Material at any time; however, doing so
+     will not terminate this Public License.
+
+  d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
+     License.
+
+
+Section 7 -- Other Terms and Conditions.
+
+  a. The Licensor shall not be bound by any additional or different
+     terms or conditions communicated by You unless expressly agreed.
+
+  b. Any arrangements, understandings, or agreements regarding the
+     Licensed Material not stated herein are separate from and
+     independent of the terms and conditions of this Public License.
+
+
+Section 8 -- Interpretation.
+
+  a. For the avoidance of doubt, this Public License does not, and
+     shall not be interpreted to, reduce, limit, restrict, or impose
+     conditions on any use of the Licensed Material that could lawfully
+     be made without permission under this Public License.
+
+  b. To the extent possible, if any provision of this Public License is
+     deemed unenforceable, it shall be automatically reformed to the
+     minimum extent necessary to make it enforceable. If the provision
+     cannot be reformed, it shall be severed from this Public License
+     without affecting the enforceability of the remaining terms and
+     conditions.
+
+  c. No term or condition of this Public License will be waived and no
+     failure to comply consented to unless expressly agreed to by the
+     Licensor.
+
+  d. Nothing in this Public License constitutes or may be interpreted
+     as a limitation upon, or waiver of, any privileges and immunities
+     that apply to the Licensor or You, including from the legal
+     processes of any jurisdiction or authority.
+
+=======================================================================
+
+Creative Commons is not a party to its public
+licenses. Notwithstanding, Creative Commons may elect to apply one of
+its public licenses to material it publishes and in those instances
+will be considered the “Licensor.” The text of the Creative Commons
+public licenses is dedicated to the public domain under the CC0 Public
+Domain Dedication. Except for the limited purpose of indicating that
+material is shared under a Creative Commons public license or as
+otherwise permitted by the Creative Commons policies published at
+creativecommons.org/policies, Creative Commons does not authorize the
+use of the trademark "Creative Commons" or any other trademark or logo
+of Creative Commons without its prior written consent including,
+without limitation, in connection with any unauthorized modifications
+to any of its public licenses or any other arrangements,
+understandings, or agreements concerning use of licensed material. For
+the avoidance of doubt, this paragraph does not form part of the
+public licenses.
+
+Creative Commons may be contacted at creativecommons.org.

README_CN.md

@@ -0,0 +1,136 @@
+# ChatTTS
+[**English**](./README.md) | [**中文简体**](./README_CN.md)
+
+ChatTTS是专门为对话场景设计的文本转语音模型，例如LLM助手对话任务。它支持英文和中文两种语言。最大的模型使用了10万小时以上的中英文数据进行训练。在HuggingFace中开源的版本为4万小时训练且未SFT的版本.
+
+如需就模型进行正式商业咨询，请发送邮件至 **[email protected]**。对于中文用户，您可以加入我们的QQ群：~~808364215 (已满)~~ 230696694 (二群) 进行讨论。同时欢迎在GitHub上提出问题。如果遇到无法使用 **[HuggingFace](https://huggingface.co/2Noise/ChatTTS)** 的情况,可以在 [modelscope](https://www.modelscope.cn/models/pzc163/chatTTS) 上进行下载. 
+
+---
+## 亮点
+1. **对话式 TTS**: ChatTTS针对对话式任务进行了优化，实现了自然流畅的语音合成，同时支持多说话人。
+2. **细粒度控制**: 该模型能够预测和控制细粒度的韵律特征，包括笑声、停顿和插入词等。
+3. **更好的韵律**: ChatTTS在韵律方面超越了大部分开源TTS模型。同时提供预训练模型，支持进一步的研究。
+
+对于模型的具体介绍, 可以参考B站的 **[宣传视频](https://www.bilibili.com/video/BV1zn4y1o7iV)**
+
+---
+
+## 免责声明
+本文件中的信息仅供学术交流使用。其目的是用于教育和研究，不得用于任何商业或法律目的。作者不保证信息的准确性、完整性或可靠性。本文件中使用的信息和数据，仅用于学术研究目的。这些数据来自公开可用的来源，作者不对数据的所有权或版权提出任何主张。
+
+ChatTTS是一个强大的文本转语音系统。然而，负责任地和符合伦理地利用这项技术是非常重要的。为了限制ChatTTS的使用，我们在4w小时模型的训练过程中添加了少量额外的高频噪音，并用mp3格式尽可能压低了音质，以防不法分子用于潜在的犯罪可能。同时我们在内部训练了检测模型，并计划在未来开放。
+
+---
+## 用法
+
+<h4>基本用法</h4>
+
+```python
+import ChatTTS
+from IPython.display import Audio
+
+chat = ChatTTS.Chat()
+chat.load_models(compile=False) # 设置为True以获得更快速度
+
+texts = ["在这里输入你的文本",]
+
+wavs = chat.infer(texts, use_decoder=True)
+
+torchaudio.save("output1.wav", torch.from_numpy(wavs[0]), 24000)
+```
+
+<h4>进阶用法</h4>
+
+```python
+###################################
+# Sample a speaker from Gaussian.
+
+rand_spk = chat.sample_random_speaker()
+
+params_infer_code = {
+  'spk_emb': rand_spk, # add sampled speaker 
+  'temperature': .3, # using custom temperature
+  'top_P': 0.7, # top P decode
+  'top_K': 20, # top K decode
+}
+
+###################################
+# For sentence level manual control.
+
+# use oral_(0-9), laugh_(0-2), break_(0-7) 
+# to generate special token in text to synthesize.
+params_refine_text = {
+  'prompt': '[oral_2][laugh_0][break_6]'
+} 
+
+wav = chat.infer(texts, params_refine_text=params_refine_text, params_infer_code=params_infer_code)
+
+###################################
+# For word level manual control.
+# use_decoder=False to infer faster with a bit worse quality
+text = 'What is [uv_break]your favorite english food?[laugh][lbreak]'
+wav = chat.infer(text, skip_refine_text=True, params_infer_code=params_infer_code, use_decoder=False)
+
+torchaudio.save("output2.wav", torch.from_numpy(wavs[0]), 24000)
+```
+
+<details open>
+  <summary><h4>自我介绍样例</h4></summary>
+
+```python
+inputs_cn = """
+chat T T S 是一款强大的对话式文本转语音模型。它有中英混读和多说话人的能力。
+chat T T S 不仅能够生成自然流畅的语音，还能控制[laugh]笑声啊[laugh]，
+停顿啊[uv_break]语气词啊等副语言现象[uv_break]。这个韵律超越了许多开源模型[uv_break]。
+请注意，chat T T S 的使用应遵守法律和伦理准则，避免滥用的安全风险。[uv_break]'
+""".replace('\n', '')
+
+params_refine_text = {
+  'prompt': '[oral_2][laugh_0][break_4]'
+} 
+audio_array_cn = chat.infer(inputs_cn, params_refine_text=params_refine_text)
+# audio_array_en = chat.infer(inputs_en, params_refine_text=params_refine_text)
+
+torchaudio.save("output3.wav", torch.from_numpy(audio_array_cn[0]), 24000)
+```
+[男说话人](https://github.com/2noise/ChatTTS/assets/130631963/bbfa3b83-2b67-4bb6-9315-64c992b63788)
+
+[女说话人](https://github.com/2noise/ChatTTS/assets/130631963/e061f230-0e05-45e6-8e4e-0189f2d260c4)
+</details>
+
+
+---
+## 计划路线
+- [x] 开源4w小时基础模型和spk_stats文件
+- [ ] 开源VQ encoder和Lora 训练代码
+- [ ] 在非refine text情况下, 流式生成音频*
+- [ ] 开源多情感可控的4w小时版本
+- [ ] ChatTTS.cpp maybe? (欢迎社区PR或独立的新repo)
+
+---
+## 常见问题
+
+##### 连不上HuggingFace
+请使用[modelscope](https://www.modelscope.cn/models/pzc163/chatTTS)的版本. 并设置cache的位置:
+```python
+chat.load_models(source='local', local_path='你的下载位置')
+```
+
+##### 我要多少显存? Infer的速度是怎么样的?
+对于30s的音频, 至少需要4G的显存. 对于4090, 1s生成约7个字所对应的音频. RTF约0.3.
+
+##### 模型稳定性似乎不够好, 会出现其他说话人或音质很差的现象.
+这是自回归模型通常都会出现的问题. 说话人可能会在中间变化, 可能会采样到音质非常差的结果, 这通常难以避免. 可以多采样几次来找到合适的结果.
+
+##### 除了笑声还能控制什么吗? 还能控制其他情感吗?
+在现在放出的模型版本中, 只有[laugh]和[uv_break], [lbreak]作为字级别的控制单元. 在未来的版本中我们可能会开源其他情感控制的版本.
+
+---
+## 致谢
+- [bark](https://github.com/suno-ai/bark),[XTTSv2](https://github.com/coqui-ai/TTS)和[valle](https://arxiv.org/abs/2301.02111)展示了自回归任务用于TTS任务的可能性.
+- [fish-speech](https://github.com/fishaudio/fish-speech)一个优秀的自回归TTS模型, 揭示了GVQ用于LLM任务的可能性.
+- [vocos](https://github.com/gemelo-ai/vocos)作为模型中的vocoder.
+
+---
+## 特别致谢
+- [wlu-audio lab](https://audio.westlake.edu.cn/)为我们提供了早期算法试验的支持.

requirements.txt

@@ -0,0 +1,8 @@
+omegaconf~=2.3.0
+torch~=2.1.0
+tqdm
+einops
+vector_quantize_pytorch
+transformers~=4.41.1
+vocos
+IPython

webui.py

@@ -0,0 +1,113 @@
+import os
+import random
+import argparse
+
+import torch
+import gradio as gr
+import numpy as np
+
+import ChatTTS
+
+
+def generate_seed():
+    new_seed = random.randint(1, 100000000)
+    return {
+        "__type__": "update",
+        "value": new_seed
+        }
+
+
+def generate_audio(text, temperature, top_P, top_K, audio_seed_input, text_seed_input, refine_text_flag):
+
+    torch.manual_seed(audio_seed_input)
+    rand_spk = chat.sample_random_speaker()
+    params_infer_code = {
+        'spk_emb': rand_spk, 
+        'temperature': temperature,
+        'top_P': top_P,
+        'top_K': top_K,
+        }
+    params_refine_text = {'prompt': '[oral_2][laugh_0][break_6]'}
+    
+    torch.manual_seed(text_seed_input)
+
+    if refine_text_flag:
+        text = chat.infer(text, 
+                          skip_refine_text=False,
+                          refine_text_only=True,
+                          params_refine_text=params_refine_text,
+                          params_infer_code=params_infer_code
+                          )
+    
+    wav = chat.infer(text, 
+                     skip_refine_text=True, 
+                     params_refine_text=params_refine_text, 
+                     params_infer_code=params_infer_code
+                     )
+    
+    audio_data = np.array(wav[0]).flatten()
+    sample_rate = 24000
+    text_data = text[0] if isinstance(text, list) else text
+
+    return [(sample_rate, audio_data), text_data]
+
+
+def main():
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# ChatTTS Webui")
+        gr.Markdown("ChatTTS Model: [2noise/ChatTTS](https://github.com/2noise/ChatTTS)")
+
+        default_text = "四川美食确实以辣闻名，但也有不辣的选择。比如甜水面、赖汤圆、蛋烘糕、叶儿粑等，这些小吃口味温和，甜而不腻，也很受欢迎。"        
+        text_input = gr.Textbox(label="Input Text", lines=4, placeholder="Please Input Text...", value=default_text)
+
+        with gr.Row():
+            refine_text_checkbox = gr.Checkbox(label="Refine text", value=True)
+            temperature_slider = gr.Slider(minimum=0.00001, maximum=1.0, step=0.00001, value=0.3, label="Audio temperature")
+            top_p_slider = gr.Slider(minimum=0.1, maximum=0.9, step=0.05, value=0.7, label="top_P")
+            top_k_slider = gr.Slider(minimum=1, maximum=20, step=1, value=20, label="top_K")
+
+        with gr.Row():
+            audio_seed_input = gr.Number(value=2, label="Audio Seed")
+            generate_audio_seed = gr.Button("\U0001F3B2")
+            text_seed_input = gr.Number(value=42, label="Text Seed")
+            generate_text_seed = gr.Button("\U0001F3B2")
+
+        generate_button = gr.Button("Generate")
+        
+        text_output = gr.Textbox(label="Output Text", interactive=False)
+        audio_output = gr.Audio(label="Output Audio")
+
+        generate_audio_seed.click(generate_seed, 
+                                  inputs=[], 
+                                  outputs=audio_seed_input)
+        
+        generate_text_seed.click(generate_seed, 
+                                 inputs=[], 
+                                 outputs=text_seed_input)
+        
+        generate_button.click(generate_audio, 
+                              inputs=[text_input, temperature_slider, top_p_slider, top_k_slider, audio_seed_input, text_seed_input, refine_text_checkbox], 
+                              outputs=[audio_output, text_output])
+
+    parser = argparse.ArgumentParser(description='ChatTTS demo Launch')
+    parser.add_argument('--server_name', type=str, default='0.0.0.0', help='Server name')
+    parser.add_argument('--server_port', type=int, default=8080, help='Server port')
+    parser.add_argument('--local_path', type=str, default=None, help='the local_path if need')
+    args = parser.parse_args()
+
+    print("loading ChatTTS model...")
+    global chat
+    chat = ChatTTS.Chat()
+
+    if args.local_path == None:
+        chat.load_models()
+    else:
+        print('local model path:', args.local_path)
+        chat.load_models('local', local_path=args.local_path)
+
+    demo.launch(server_name=args.server_name, server_port=args.server_port, inbrowser=True)
+
+
+if __name__ == '__main__':
+    main()