webui origin

ChatTTS/__init__.py

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

ChatTTS/core.py

@@ -0,0 +1,149 @@
+
+import os
+import logging
+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.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.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>'):
+        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}')
+            self._load(**{k: os.path.join(download_path, v) for k, v in OmegaConf.load(os.path.join(download_path, 'config', 'path.yaml')).items()})
+        elif source == 'local':
+            self.logger.log(logging.INFO, f'Load from local: {local_path}')
+            self._load(**{k: os.path.join(local_path, v) for k, v in OmegaConf.load(os.path.join(local_path, 'config', 'path.yaml')).items()})
+        
+    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
+    ):
+        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'))
+            self.pretrain_models['gpt'] = gpt
+            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={}, 
+        use_decoder=False
+    ):
+        
+        assert self.check_model(use_decoder=use_decoder)
+        
+        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
+        
+        
+

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}[uv_break][Ptts]' for i in text] 
+    else:
+        text = [f'[Stts][empty_spk]{i}[uv_break][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,45 @@
+
+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

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

app.py

@@ -0,0 +1,105 @@
+import os
+import random
+import argparse
+
+import torch
+import gradio as gr
+import numpy as np
+
+import ChatTTS
+
+print("loading ChatTTS model...")
+chat = ChatTTS.Chat()
+chat.load_models()
+
+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 = torch.randn(768)
+    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=42, 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')
+    args = parser.parse_args()
+
+    demo.launch(server_name=args.server_name, server_port=args.server_port, inbrowser=True)
+
+
+if __name__ == '__main__':
+    main()

download_model.py

@@ -0,0 +1,5 @@
+#SDK模型下载
+from modelscope import snapshot_download
+model_dir = snapshot_download('pzc163/chatTTS')
+
+print(model_dir)

main.ipynb

@@ -0,0 +1,27 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import ChatTTS\n",
+    "from IPython.display import Audio"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "chat-tts",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

main.py

@@ -0,0 +1,10 @@
+import ChatTTS
+from IPython.display import Audio
+
+chat = ChatTTS.Chat()
+chat.load_models(source = "local", local_path='/home/u210110703/.cache/modelscope/hub/pzc163/chatTTS')
+
+texts = ["我觉得像我们这些写程序的人，他，我觉得多多少少可能会对开源有一种情怀在吧我觉得开源是一个很好的形式。现在其实最先进的技术掌握在一些公司的手里的话，就他们并不会轻易的开放给所有的人用。",]
+
+wavs = chat.infer(texts, use_decoder=True)
+Audio(wavs[0], rate=24_000, autoplay=True)

webui.py

@@ -0,0 +1,105 @@
+import os
+import random
+import argparse
+
+import torch
+import gradio as gr
+import numpy as np
+
+import ChatTTS
+
+print("loading ChatTTS model...")
+chat = ChatTTS.Chat()
+chat.load_models()
+
+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 = torch.randn(768)
+    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=42, 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')
+    args = parser.parse_args()
+
+    demo.launch(server_name=args.server_name, server_port=args.server_port, inbrowser=True)
+
+
+if __name__ == '__main__':
+    main()