新增gpu模式下chattts代码

Chat2TTS/__init__.py

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

Chat2TTS/core.py

@@ -0,0 +1,171 @@
+
+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 dataclasses import dataclass
+from typing import Literal, Optional, List, Tuple, Dict
+
+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()
+
+    @dataclass(repr=False, eq=False)
+    class RefineTextParams:
+        prompt: str = ""
+        top_P: float = 0.7
+        top_K: int = 20
+        temperature: float = 0.7
+        repetition_penalty: float = 1.0
+        max_new_token: int = 384
+        min_new_token: int = 0
+        show_tqdm: bool = True
+        ensure_non_empty: bool = True
+
+    @dataclass(repr=False, eq=False)
+    class InferCodeParams(RefineTextParams):
+        prompt: str = "[speed_5]"
+        spk_emb: Optional[str] = None
+        temperature: float = 0.3
+        repetition_penalty: float = 1.05
+        max_new_token: int = 2048
+
+    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
+        
+        
+

Chat2TTS/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

Chat2TTS/infer/api.py

@@ -0,0 +1,126 @@
+
+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

Chat2TTS/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

Chat2TTS/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,
+            }

Chat2TTS/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

Chat2TTS/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

Chat2TTS/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

tool/gpu.py

@@ -1,12 +1,11 @@
 import torch
 import os, sys
-import spaces
 
 if sys.platform == "darwin":
     os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
 now_dir = os.getcwd()
 sys.path.append(now_dir)
-from tool.logger import get_logger
+from .logger.log import get_logger
 
 logger = get_logger("gpu")
 
@@ -36,13 +35,13 @@ def select_device(min_memory=2047, experimental=False):
         """
         if experimental:
             # For Apple M1/M2 chips with Metal Performance Shaders
-            logger.get_logger().warn("experimantal: found apple GPU, using MPS.")
+            logger.warn("experimantal: found apple GPU, using MPS.")
             device = torch.device("mps")
         else:
-            logger.get_logger().info("found Apple GPU, but use CPU.")
+            logger.info("found Apple GPU, but use CPU.")
             device = torch.device("cpu")
     else:
-        logger.get_logger().warning("no GPU found, use CPU instead")
+        logger.warning("no GPU found, use CPU instead")
         device = torch.device("cpu")
 
     return device

web/app.py

@@ -10,7 +10,7 @@ from tool.logger import get_logger
 from tool.func import *
 from tool.np import *
 from tool.gpu import select_device
-import ChatTTS
+import Chat2TTS
 import argparse
 import torch._dynamo
 
@@ -21,7 +21,7 @@ torch._dynamo.config.suppress_errors = True
 logger = get_logger("app")
 
 # Initialize and load the model:
-chat = ChatTTS.Chat()
+chat = Chat2TTS.Chat()
 
 
 def init_chat(args):
@@ -38,7 +38,7 @@ def init_chat(args):
 
     logger.info("loading ChatTTS device :" + str(device))
 
-    if chat.load(source=source, custom_path="D:\\chenjgspace\\ai-model\\chattts", device=device):
+    if chat.load_models(source=source, local_path="D:\\chenjgspace\\ai-model\\chattts"):
         print("Models loaded successfully.")
         logger.info("Models loaded successfully.")
     else:
@@ -203,7 +203,7 @@ def get_chat_infer_audio(chat_txt,
                  spk_emb_text):
     logger.info("========开始生成音频文件=====")
     #音频参数设置
-    params_infer_code = ChatTTS.Chat.InferCodeParams(
+    params_infer_code = Chat2TTS.Chat.InferCodeParams(
         spk_emb=spk_emb_text,  # add sampled speaker
         temperature=temperature_slider,  # using custom temperature
         top_P=top_p_slider,  # top P decode
@@ -227,7 +227,7 @@ def get_chat_infer_text(text,seed,refine_text_checkBox):
         logger.info("========文本内容无需优化=====")
         return  text
 
-    params_refine_text = ChatTTS.Chat.RefineTextParams(
+    params_refine_text = Chat2TTS.Chat.RefineTextParams(
         prompt='[oral_2][laugh_0][break_6]',
     )
 

web/app_sc.py

@@ -0,0 +1,267 @@
+import os, sys
+import spaces
+
+if sys.platform == "darwin":
+    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+
+from tool.logger import get_logger
+from tool.func import *
+from tool.np import *
+from tool.gpu import select_device
+from tool.ctx import TorchSeedContext
+import ChatTTS
+import argparse
+import torch._dynamo
+
+torch._dynamo.config.suppress_errors = True
+
+
+
+logger = get_logger("app")
+
+# Initialize and load the model:
+chat = ChatTTS.Chat()
+
+
+def init_chat(args):
+    global chat
+    source = "custom"
+    # 获取启动模式
+    MODEL = os.getenv('MODEL')
+    logger.info("loading ChatTTS model..., start MODEL:" + str(MODEL))
+    # huggingface 部署模式下，模型则直接使用hf的模型数据
+    if MODEL == "HF":
+        source = "huggingface"
+
+    device=select_device()
+
+    logger.info("loading ChatTTS device :" + str(device))
+
+    if chat.load(source=source, custom_path="D:\\chenjgspace\\ai-model\\chattts",device=device):
+        print("Models loaded successfully.")
+        logger.info("Models loaded successfully.")
+    else:
+        logger.error("=========Models load failed.")
+        sys.exit(1)
+
+
+def main(args):
+    with gr.Blocks() as demo:
+        gr.Markdown("# ChatTTS demo")
+        with gr.Row():
+            with gr.Column(scale=1):
+                text_input = gr.Textbox(
+                    label="转换内容",
+                    lines=4,
+                    max_lines=4,
+                    placeholder="Please Input Text...",
+                    value="柔柔的，浓浓的，痴痴的风，牵引起心底灵动的思潮；情愫悠悠，思情绵绵，风里默坐，红尘中的浅醉，诗词中的优柔，任那自在飞花轻似梦的情怀，裁一束霓衣，织就清浅淡薄的安寂。",
+                    interactive=True,
+                )
+        with gr.Row():
+            refine_text_checkBox = gr.Checkbox(
+                label="是否优化文本,如是则先对文本内容做优化分词",
+                interactive=True,
+                value=True
+            )
+            temperature_slider = gr.Slider(
+                minimum=0.00001,
+                maximum=1.0,
+                step=0.00001,
+                value=0.3,
+                interactive=True,
+                label="模型 Temperature 参数设置"
+            )
+            top_p_slider = gr.Slider(
+                minimum=0.1,
+                maximum=0.9,
+                step=0.05,
+                value=0.7,
+                label="模型 top_P 参数设置",
+                interactive=True,
+            )
+            top_k_slider = gr.Slider(
+                minimum=1,
+                maximum=20,
+                step=1,
+                value=20,
+                label="模型 top_K 参数设置",
+                interactive=True,
+            )
+        with gr.Row():
+            voice_selection = gr.Dropdown(
+                label="Timbre",
+                choices=voices.keys(),
+                value="旁白",
+                interactive=True,
+                show_label=True
+            )
+            audio_seed_input = gr.Number(
+                value=2,
+                label="音色种子",
+                interactive=True,
+                minimum=seed_min,
+                maximum=seed_max,
+            )
+            generate_audio_seed = gr.Button("随机生成音色种子", interactive=True)
+            text_seed_input = gr.Number(
+                value=42,
+                label="文本种子",
+                interactive=True,
+                minimum=seed_min,
+                maximum=seed_max,
+            )
+            generate_text_seed = gr.Button("随机生成文本种子", interactive=True)
+
+        with gr.Row():
+            spk_emb_text = gr.Textbox(
+                label="Speaker Embedding",
+                max_lines=3,
+                show_copy_button=True,
+                interactive=False,
+                scale=2,
+
+            )
+            reload_chat_button = gr.Button("Reload", scale=1, interactive=True)
+
+        with gr.Row():
+            generate_button = gr.Button("生成音频文件", scale=1, interactive=True)
+
+        with gr.Row():
+            text_output = gr.Textbox(
+                label="输出文本",
+                interactive=False,
+                show_copy_button=True,
+            )
+
+            audio_output = gr.Audio(
+                label="输出音频",
+                value=None,
+                format="wav",
+                autoplay=False,
+                streaming=False,
+                interactive=False,
+                show_label=True,
+                waveform_options=gr.WaveformOptions(
+                    sample_rate=24000,
+                ),
+            )
+        # 针对页面元素新增 监听事件
+        voice_selection.change(fn=on_voice_change, inputs=voice_selection, outputs=audio_seed_input)
+
+        audio_seed_input.change(fn=on_audio_seed_change, inputs=audio_seed_input, outputs=spk_emb_text)
+
+        generate_audio_seed.click(fn=generate_seed, outputs=audio_seed_input)
+
+        generate_text_seed.click(fn=generate_seed,outputs=text_seed_input)
+
+        # reload_chat_button.click()
+
+        generate_button.click(fn=get_chat_infer_text,
+                              inputs=[text_input,
+                                                       text_seed_input,
+                                                       refine_text_checkBox
+                                                       ],
+                              outputs=[text_output]
+                              ).then(fn=get_chat_infer_audio,
+                                     inputs=[text_output,
+                                                       temperature_slider,
+                                                       top_p_slider,
+                                                       top_k_slider,
+                                                       audio_seed_input,
+                                                       spk_emb_text
+                                                       ],
+                                     outputs=[audio_output])
+        # 初始化 spk_emb_text 数值
+        spk_emb_text.value = on_audio_seed_change(audio_seed_input.value)
+        logger.info("元素初始化完成,启动gradio服务=======")
+
+        # 运行gradio服务
+        demo.launch()
+
+
+    '''
+    top_K: "top_K"（K个最高得分）是指在所有可能的生成结果中，模型会选取前K个得分最高的结果。这个设置常常用于基于概率的生成任务，例如语言模型中的词或句子生成。当你设置top_K为K时，你要求模型选择得分最高的K个选项，这样输出通常会有一定的多样性，但仍然是基于模型的前K个预测。
+
+    top_P: "top_P"（概率阈值）则是一个连续的值，而不是离散的整数。它代表的是一个概率阈值，模型会生成所有得分高于该阈值的概率的项目。换句话说，top_P会生成那些概率大于等于给定值的所有生成结果。这个设置更为灵活，可以根据实际需求调整生成内容的不确定性，高频选项被生成的概率较高，而低频可能性则可能根据阈值随机出现。
+
+在实际应用中，选择top_K还是top_P取决于具体任务需求，如是否希望生成内容有一定程度的多样化（top_K），还是希望生成的内容更接近于最可能发生的选项（top_P）。较高的top_P可能会引入更多的随机性和创新，而较低的top_K则更倾向于保守的选择。
+
+    spk_embedding（Speaker Embedding）: 这个术语一般用于语音识别或者多说话者模型中。"spk_embedding"指的是每个说话人的身份或特征向量，或者说是用户标识的嵌入表示。在对话系统中，它能帮助模型区分不同的说话者，比如在多轮对话中区分是同一个用户的不同回复，或者是不同用户的交互。这个嵌入可能包含了说话人的个性、语调、口音等信息，有助于提高对话的连贯性和自然性。
+
+    temperature: 通常在语言模型的生成（如基于概率的softmax）中使用。"temperature"是一个正数，用于控制生成内容的随机性和多样性。当温度较低（如0.1）时，模型倾向于生成最可能的结果，文字更保守，少有创新；当温度较高（如1或更高）时，模型将更倾向于产生多样化的内容，但可能性较大的选项将被稀释。因此，temperature调整是一个常用的平衡方法，使得生成更具创造性或是更符合预期。
+
+简而言之，"spk_embedding"关注的是对话参与者的身份特征，而"temperature"是用于调整生成文本不确定性的一个超参数。
+    '''
+#@spaces.GPU  使用chattts三方包 gpu模式没有跑通
+def get_chat_infer_audio(chat_txt,
+                 temperature_slider,
+                 top_p_slider,
+                 top_k_slider,
+                 audio_seed_input,
+                 spk_emb_text):
+    logger.info("========开始生成音频文件=====")
+    #音频参数设置
+    params_infer_code = ChatTTS.Chat.InferCodeParams(
+        spk_emb=spk_emb_text,  # add sampled speaker
+        temperature=temperature_slider,  # using custom temperature
+        top_P=top_p_slider,  # top P decode
+        top_K=top_k_slider,  # top K decode
+    )
+
+    with TorchSeedContext(audio_seed_input):
+        wav = chat.infer(
+            text=chat_txt,
+            skip_refine_text=True, #跳过文本优化
+            params_infer_code=params_infer_code,
+        )
+    yield 24000, float_to_int16(wav[0]).T
+
+#@spaces.GPU  使用chattts三方包 gpu模式没有跑通
+def get_chat_infer_text(text,seed,refine_text_checkBox):
+
+    logger.info("========开始优化文本内容2=====")
+    global chat
+    if not refine_text_checkBox:
+        logger.info("========文本内容无需优化=====")
+        return  text
+
+    params_refine_text = ChatTTS.Chat.RefineTextParams(
+        prompt='[oral_2][laugh_0][break_6]',
+    )
+
+    with TorchSeedContext(seed):
+        chat_text = chat.infer(
+            text=text,
+            skip_refine_text=False,
+            refine_text_only=True,  #仅返回优化后文本内容
+            params_refine_text=params_refine_text,
+        )
+
+    return chat_text[0] if isinstance(chat_text, list) else chat_text
+
+@spaces.GPU
+def on_audio_seed_change(audio_seed_input):
+    global chat
+    with TorchSeedContext(audio_seed_input):
+        rand_spk = chat.sample_random_speaker()
+    return rand_spk
+
+
+if __name__ == "__main__":
+    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=7860, help="server port")
+    parser.add_argument(
+        "--custom_path", type=str, default="D:\\chenjgspace\\ai-model\\chattts", help="custom model path"
+    )
+    parser.add_argument(
+        "--coef", type=str, default=None, help="custom dvae coefficient"
+    )
+    args = parser.parse_args()
+    init_chat(args)
+    main(args)