diff --git a/ChatTTS/core.py b/ChatTTS/core.py index 050f719..9359837 100644 --- a/ChatTTS/core.py +++ b/ChatTTS/core.py @@ -15,7 +15,13 @@ from huggingface_hub import snapshot_download from transformers.generation import TopKLogitsWarper, TopPLogitsWarper from .model import DVAE, GPT, CustomRepetitionPenaltyLogitsProcessorRepeat -from .utils import check_all_assets, download_all_assets, select_device, get_latest_modified_file, del_all +from .utils import ( + check_all_assets, + download_all_assets, + select_device, + get_latest_modified_file, + del_all, +) from .utils import logger as utils_logger from .norm import Normalizer @@ -28,76 +34,89 @@ class Chat: self.pretrain_models = {} self.normalizer = Normalizer( - os.path.join(os.path.dirname(__file__), 'res', 'homophones_map.json'), + os.path.join(os.path.dirname(__file__), "res", "homophones_map.json"), logger, ) - with open(os.path.join(os.path.dirname(__file__), 'res', 'sha256_map.json')) as f: + with open( + os.path.join(os.path.dirname(__file__), "res", "sha256_map.json") + ) as f: self.sha256_map: Dict[str, str] = load(f) self.context = GPT.Context() - def has_loaded(self, use_decoder = False): + def has_loaded(self, use_decoder=False): not_finish = False - check_list = ["vocos", "_vocos_decode", 'gpt', 'tokenizer'] - + check_list = ["vocos", "_vocos_decode", "gpt", "tokenizer"] + if use_decoder: - check_list.append('decoder') + check_list.append("decoder") else: - check_list.append('dvae') + check_list.append("dvae") for module in check_list: if not hasattr(self, module) and module not in self.pretrain_models: - self.logger.warning(f'{module} not initialized.') + self.logger.warning(f"{module} not initialized.") not_finish = True if not not_finish: - self.logger.info('all models has been initialized.') + self.logger.info("all models has been initialized.") return not not_finish def download_models( self, - source: Literal['huggingface', 'local', 'custom']='local', + source: Literal["huggingface", "local", "custom"] = "local", force_redownload=False, - custom_path: Optional[torch.serialization.FILE_LIKE]=None, + custom_path: Optional[torch.serialization.FILE_LIKE] = None, ) -> Optional[str]: - if source == 'local': + if source == "local": download_path = os.getcwd() if not check_all_assets(self.sha256_map, update=True) or force_redownload: with tempfile.TemporaryDirectory() as tmp: download_all_assets(tmpdir=tmp) if not check_all_assets(self.sha256_map, update=False): - self.logger.error("download to local path %s failed.", download_path) + self.logger.error( + "download to local path %s failed.", download_path + ) return None - elif source == 'huggingface': - hf_home = os.getenv('HF_HOME', os.path.expanduser("~/.cache/huggingface")) + elif 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')) + 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') + if download_path is None or force_redownload: + self.logger.log( + logging.INFO, + f"download from HF: https://huggingface.co/2Noise/ChatTTS", + ) try: - download_path = snapshot_download(repo_id="2Noise/ChatTTS", allow_patterns=["*.pt", "*.yaml"]) + download_path = snapshot_download( + repo_id="2Noise/ChatTTS", allow_patterns=["*.pt", "*.yaml"] + ) except: download_path = None else: - self.logger.log(logging.INFO, f'load latest snapshot from cache: {download_path}') + self.logger.log( + logging.INFO, f"load latest snapshot from cache: {download_path}" + ) if download_path is None: self.logger.error("download from huggingface failed.") return None - elif source == 'custom': - self.logger.log(logging.INFO, f'try to load from local: {custom_path}') + elif source == "custom": + self.logger.log(logging.INFO, f"try to load from local: {custom_path}") download_path = custom_path - + return download_path def load( self, - source: Literal['huggingface', 'local', 'custom']='local', + source: Literal["huggingface", "local", "custom"] = "local", force_redownload=False, compile: bool = True, - custom_path: Optional[torch.serialization.FILE_LIKE]=None, + custom_path: Optional[torch.serialization.FILE_LIKE] = None, device: Optional[torch.device] = None, coef: Optional[torch.Tensor] = None, ) -> bool: @@ -105,10 +124,17 @@ class Chat: if download_path is None: return False return self._load( - device=device, compile=compile, coef=coef, - **{k: os.path.join(download_path, v) for k, v in OmegaConf.load(os.path.join(download_path, 'config', 'path.yaml')).items()}, + device=device, + compile=compile, + coef=coef, + **{ + k: os.path.join(download_path, v) + for k, v in OmegaConf.load( + os.path.join(download_path, "config", "path.yaml") + ).items() + }, ) - + def unload(self): logger = self.logger del_all(self.pretrain_models) @@ -116,7 +142,7 @@ class Chat: del self.normalizer del self.sha256_map self._gen_logits.cache_clear() - del_list = ["vocos", "_vocos_decode", 'gpt', 'decoder', 'dvae'] + del_list = ["vocos", "_vocos_decode", "gpt", "decoder", "dvae"] for module in del_list: if hasattr(self, module): delattr(self, module) @@ -124,12 +150,12 @@ class Chat: def sample_random_speaker(self): dim = self.gpt.gpt.layers[0].mlp.gate_proj.in_features - std, mean = self.pretrain_models['spk_stat'].chunk(2) + std, mean = self.pretrain_models["spk_stat"].chunk(2) return torch.randn(dim, device=std.device) * std + mean @dataclass(repr=False, eq=False) - class RefineTextParams(): - prompt: str = '' + class RefineTextParams: + prompt: str = "" top_P: float = 0.7 top_K: int = 20 temperature: float = 0.7 @@ -138,8 +164,8 @@ class Chat: min_new_token: int = 0 @dataclass(repr=False, eq=False) - class InferCodeParams(): - prompt: str = '[speed_5]' + class InferCodeParams: + prompt: str = "[speed_5]" spk_emb: Optional[torch.Tensor] = None top_P: float = 0.7 top_K: int = 20 @@ -149,42 +175,42 @@ class Chat: min_new_token: int = 0 def infer( - self, + self, text, stream=False, lang=None, - skip_refine_text=False, + skip_refine_text=False, refine_text_only=False, use_decoder=True, do_text_normalization=True, do_homophone_replacement=True, - params_refine_text = RefineTextParams(), - params_infer_code = InferCodeParams(), + params_refine_text=RefineTextParams(), + params_infer_code=InferCodeParams(), ): self.context.set(False) res_gen = self._infer( text, stream, lang, - skip_refine_text, + skip_refine_text, refine_text_only, use_decoder, do_text_normalization, do_homophone_replacement, - params_refine_text, - params_infer_code, + params_refine_text, + params_infer_code, ) if stream: return res_gen else: return next(res_gen) - + def interrupt(self): self.context.set(True) def _load( - self, - vocos_config_path: str = None, + self, + vocos_config_path: str = None, vocos_ckpt_path: str = None, dvae_config_path: str = None, dvae_ckpt_path: str = None, @@ -195,105 +221,141 @@ class Chat: tokenizer_path: str = None, device: Optional[torch.device] = None, compile: bool = True, - coef: Optional[str] = None + coef: Optional[str] = None, ): if device is None: device = select_device() - self.logger.log(logging.INFO, f'use {device}') + self.logger.log(logging.INFO, f"use {device}") self.device = device if vocos_config_path: - vocos = Vocos.from_hparams(vocos_config_path).to( - # vocos on mps will crash, use cpu fallback - "cpu" if "mps" in str(device) else device - ).eval() - assert vocos_ckpt_path, 'vocos_ckpt_path should not be None' - vocos.load_state_dict(torch.load(vocos_ckpt_path, weights_only=True, mmap=True)) + vocos = ( + Vocos.from_hparams(vocos_config_path) + .to( + # vocos on mps will crash, use cpu fallback + "cpu" + if "mps" in str(device) + else device + ) + .eval() + ) + assert vocos_ckpt_path, "vocos_ckpt_path should not be None" + vocos.load_state_dict( + torch.load(vocos_ckpt_path, weights_only=True, mmap=True) + ) self.vocos = vocos if "mps" in str(self.device): - self._vocos_decode: Callable[[torch.Tensor], np.ndarray] = lambda spec: self.vocos.decode( - spec.cpu() - ).cpu().numpy() + self._vocos_decode: Callable[[torch.Tensor], np.ndarray] = ( + lambda spec: self.vocos.decode(spec.cpu()).cpu().numpy() + ) else: - self._vocos_decode: Callable[[torch.Tensor], np.ndarray] = lambda spec: self.vocos.decode( - spec - ).cpu().numpy() - self.logger.log(logging.INFO, 'vocos loaded.') + self._vocos_decode: Callable[[torch.Tensor], np.ndarray] = ( + lambda spec: self.vocos.decode(spec).cpu().numpy() + ) + self.logger.log(logging.INFO, "vocos loaded.") if dvae_config_path: cfg = OmegaConf.load(dvae_config_path) dvae = DVAE(**cfg, coef=coef).to(device).eval() coef = str(dvae) - assert dvae_ckpt_path, 'dvae_ckpt_path should not be None' - dvae.load_state_dict(torch.load(dvae_ckpt_path, weights_only=True, mmap=True)) + assert dvae_ckpt_path, "dvae_ckpt_path should not be None" + dvae.load_state_dict( + torch.load(dvae_ckpt_path, weights_only=True, mmap=True) + ) self.dvae = dvae - self.logger.log(logging.INFO, 'dvae loaded.') - + self.logger.log(logging.INFO, "dvae loaded.") + if gpt_config_path: cfg = OmegaConf.load(gpt_config_path) gpt = GPT(**cfg, device=device, logger=self.logger).eval() - assert gpt_ckpt_path, 'gpt_ckpt_path should not be None' + assert gpt_ckpt_path, "gpt_ckpt_path should not be None" gpt.load_state_dict(torch.load(gpt_ckpt_path, weights_only=True, mmap=True)) - if compile and 'cuda' in str(device): + if compile and "cuda" in str(device): try: - gpt.gpt.forward = torch.compile(gpt.gpt.forward, backend='inductor', dynamic=True) + gpt.gpt.forward = torch.compile( + gpt.gpt.forward, backend="inductor", dynamic=True + ) except RuntimeError as e: - self.logger.warning(f'compile failed: {e}. fallback to normal mode.') + self.logger.warning( + f"compile failed: {e}. fallback to normal mode." + ) self.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, weights_only=True, mmap=True).to(device) - self.logger.log(logging.INFO, 'gpt loaded.') - + 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, weights_only=True, mmap=True + ).to(device) + self.logger.log(logging.INFO, "gpt loaded.") + if decoder_config_path: cfg = OmegaConf.load(decoder_config_path) decoder = DVAE(**cfg, coef=coef).to(device).eval() coef = str(decoder) - assert decoder_ckpt_path, 'decoder_ckpt_path should not be None' - decoder.load_state_dict(torch.load(decoder_ckpt_path, weights_only=True, mmap=True)) + assert decoder_ckpt_path, "decoder_ckpt_path should not be None" + decoder.load_state_dict( + torch.load(decoder_ckpt_path, weights_only=True, mmap=True) + ) self.decoder = decoder - self.logger.log(logging.INFO, 'decoder loaded.') - + self.logger.log(logging.INFO, "decoder loaded.") + if tokenizer_path: tokenizer = torch.load(tokenizer_path, map_location=device, mmap=True) - tokenizer.padding_side = 'left' - self.pretrain_models['tokenizer'] = tokenizer - self.logger.log(logging.INFO, 'tokenizer loaded.') - + tokenizer.padding_side = "left" + self.pretrain_models["tokenizer"] = tokenizer + self.logger.log(logging.INFO, "tokenizer loaded.") + self.coef = coef return self.has_loaded() def _infer( - self, + self, text, stream=False, lang=None, - skip_refine_text=False, + skip_refine_text=False, refine_text_only=False, use_decoder=True, do_text_normalization=True, do_homophone_replacement=True, - params_refine_text = RefineTextParams(), - params_infer_code = InferCodeParams(), + params_refine_text=RefineTextParams(), + params_infer_code=InferCodeParams(), ): assert self.has_loaded(use_decoder=use_decoder) - if not isinstance(text, list): + if not isinstance(text, list): text = [text] - text = [self.normalizer( - t, do_text_normalization, do_homophone_replacement, lang, - ) for t in text] + text = [ + self.normalizer( + t, + do_text_normalization, + do_homophone_replacement, + lang, + ) + for t in text + ] if not skip_refine_text: refined = self._refine_text( - text, self.device, params_refine_text, + text, + self.device, + params_refine_text, ) text_tokens = refined.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) + 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) refined.destroy() if refine_text_only: yield text @@ -301,12 +363,18 @@ class Chat: length = [0 for _ in range(len(text))] for result in self._infer_code( - text, stream, self.device, use_decoder, params_infer_code, + text, + stream, + self.device, + use_decoder, + params_infer_code, ): wav = self._decode_to_wavs(result, length, use_decoder) yield wav - def _decode_to_wavs(self, result: GPT.GenerationOutputs, start_seeks: List[int], use_decoder: bool): + def _decode_to_wavs( + self, result: GPT.GenerationOutputs, start_seeks: List[int], use_decoder: bool + ): x = result.hiddens if use_decoder else result.ids wavs: List[np.ndarray] = [] for i, chunk_data in enumerate(x): @@ -318,24 +386,26 @@ class Chat: start_seeks[i] = length chunk_data = chunk_data[start_seek:] decoder = self.decoder if use_decoder else self.dvae - mel_spec = decoder(chunk_data[None].permute(0,2,1).to(self.device)) + mel_spec = decoder(chunk_data[None].permute(0, 2, 1).to(self.device)) del chunk_data wavs.append(self._vocos_decode(mel_spec)) del_all(mel_spec) result.destroy() del_all(x) return wavs - + def _gen_gpt_inputs(self, text: str, device="cpu"): gpt = self.gpt - tokenizer = self.pretrain_models['tokenizer'] + tokenizer = self.pretrain_models["tokenizer"] - text_token_tmp = tokenizer(text, return_tensors='pt', add_special_tokens=False, padding=True) + text_token_tmp = tokenizer( + text, return_tensors="pt", add_special_tokens=False, padding=True + ) text_token = text_token_tmp.to(device) del text_token_tmp - input_ids = text_token['input_ids'][...,None].expand(-1, -1, gpt.num_vq) - text_mask = torch.ones(text_token['input_ids'].shape, dtype=bool, device=device) + input_ids = text_token["input_ids"][..., None].expand(-1, -1, gpt.num_vq) + text_mask = torch.ones(text_token["input_ids"].shape, dtype=bool, device=device) return input_ids, text_token, text_mask @@ -343,9 +413,9 @@ class Chat: def _gen_logits( self, num_code: int, - top_P = 0.7, - top_K = 20, - repetition_penalty = 1.0, + top_P=0.7, + top_K=20, + repetition_penalty=1.0, ): logits_warpers = [] if top_P is not None: @@ -355,11 +425,14 @@ class Chat: logits_processors = [] if repetition_penalty is not None and repetition_penalty != 1: - logits_processors.append(CustomRepetitionPenaltyLogitsProcessorRepeat(\ - repetition_penalty, num_code, 16)) - + logits_processors.append( + CustomRepetitionPenaltyLogitsProcessorRepeat( + repetition_penalty, num_code, 16 + ) + ) + return logits_warpers, logits_processors - + def _apply_spk_emb( self, emb: torch.Tensor, @@ -368,10 +441,12 @@ class Chat: text_len: int, ): - tokenizer = self.pretrain_models['tokenizer'] + tokenizer = self.pretrain_models["tokenizer"] - n = F.normalize(spk_emb.to(emb.dtype)[None].expand(text_len, -1), p=2.0, dim=1, eps=1e-12).to(self.gpt.device_gpt) - emb[input_ids[..., 0] == tokenizer.convert_tokens_to_ids('[spk_emb]')] = n + n = F.normalize( + spk_emb.to(emb.dtype)[None].expand(text_len, -1), p=2.0, dim=1, eps=1e-12 + ).to(self.gpt.device_gpt) + emb[input_ids[..., 0] == tokenizer.convert_tokens_to_ids("[spk_emb]")] = n del n def _infer_code( @@ -385,23 +460,23 @@ class Chat: gpt = self.gpt - if not isinstance(text, list): + if not isinstance(text, list): text = [text] - - assert len(text), 'text should not be empty' + + assert len(text), "text should not be empty" if not isinstance(params.temperature, list): temperature = [params.temperature] * gpt.num_vq else: temperature = params.temperature - + if params.prompt: text = [params.prompt + i for i in text] if params.spk_emb is not None: - text = [f'[Stts][spk_emb]{i}[Ptts]' for i in text] + text = [f"[Stts][spk_emb]{i}[Ptts]" for i in text] else: - text = [f'[Stts][empty_spk]{i}[Ptts]' for i in text] + text = [f"[Stts][empty_spk]{i}[Ptts]" for i in text] input_ids, text_token, text_mask = self._gen_gpt_inputs(text, gpt.device_gpt) @@ -419,19 +494,20 @@ class Chat: top_K=params.top_K, repetition_penalty=params.repetition_penalty, ) - + result = gpt.generate( - emb, input_ids, - temperature = torch.tensor(temperature, device=device), - eos_token = num_code, - attention_mask = text_token['attention_mask'], - max_new_token = params.max_new_token, - min_new_token = params.min_new_token, - logits_warpers = logits_warpers, - logits_processors = logits_processors, - infer_text = False, + emb, + input_ids, + temperature=torch.tensor(temperature, device=device), + eos_token=num_code, + attention_mask=text_token["attention_mask"], + max_new_token=params.max_new_token, + min_new_token=params.min_new_token, + logits_warpers=logits_warpers, + logits_processors=logits_processors, + infer_text=False, return_hidden=return_hidden, - stream = stream, + stream=stream, context=self.context, ) @@ -448,15 +524,15 @@ class Chat: ): gpt = self.gpt - tokenizer = self.pretrain_models['tokenizer'] + tokenizer = self.pretrain_models["tokenizer"] - if not isinstance(text, list): + if not isinstance(text, list): text = [text] text = [f"[Sbreak]{i}[Pbreak]{params.prompt}" for i in text] input_ids, text_token, text_mask = self._gen_gpt_inputs(text, gpt.device_gpt) - + logits_warpers, logits_processors = self._gen_logits( num_code=len(tokenizer), top_P=params.top_P, @@ -468,16 +544,19 @@ class Chat: del text_mask result = gpt.generate( - emb, input_ids, - temperature = torch.tensor([params.temperature], device=device), - eos_token = torch.tensor(tokenizer.convert_tokens_to_ids('[Ebreak]'), device=gpt.device_gpt)[None], - attention_mask = text_token['attention_mask'], - max_new_token = params.max_new_token, + emb, + input_ids, + temperature=torch.tensor([params.temperature], device=device), + eos_token=torch.tensor( + tokenizer.convert_tokens_to_ids("[Ebreak]"), device=gpt.device_gpt + )[None], + attention_mask=text_token["attention_mask"], + max_new_token=params.max_new_token, min_new_token=params.min_new_token, - logits_warpers = logits_warpers, - logits_processors = logits_processors, - infer_text = True, - stream = False, + logits_warpers=logits_warpers, + logits_processors=logits_processors, + infer_text=True, + stream=False, context=self.context, ) diff --git a/ChatTTS/model/dvae.py b/ChatTTS/model/dvae.py index 4090d73..491477e 100644 --- a/ChatTTS/model/dvae.py +++ b/ChatTTS/model/dvae.py @@ -8,23 +8,31 @@ import torch.nn as nn import torch.nn.functional as F from vector_quantize_pytorch import GroupedResidualFSQ + class ConvNeXtBlock(nn.Module): def __init__( self, dim: int, intermediate_dim: int, - kernel: int, dilation: int, + kernel: int, + dilation: int, 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.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.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 = ( @@ -33,7 +41,7 @@ class ConvNeXtBlock(nn.Module): else None ) - def forward(self, x: torch.Tensor, cond = None) -> torch.Tensor: + def forward(self, x: torch.Tensor, cond=None) -> torch.Tensor: residual = x y = self.dwconv(x) @@ -58,9 +66,9 @@ class ConvNeXtBlock(nn.Module): class GFSQ(nn.Module): - def __init__(self, - dim: int, levels: List[int], G: int, R: int, eps=1e-5, transpose = True - ): + def __init__( + self, dim: int, levels: List[int], G: int, R: int, eps=1e-5, transpose=True + ): super(GFSQ, self).__init__() self.quantizer = GroupedResidualFSQ( dim=dim, @@ -73,7 +81,7 @@ class GFSQ(nn.Module): self.transpose = transpose self.G = G self.R = R - + def _embed(self, x: torch.Tensor): if self.transpose: x = x.transpose(1, 2) @@ -84,7 +92,7 @@ class GFSQ(nn.Module): """ x = x.view(x.size(0), x.size(1), self.G, self.R).permute(2, 0, 1, 3) feat = self.quantizer.get_output_from_indices(x) - return feat.transpose_(1,2) if self.transpose else feat + return feat.transpose_(1, 2) if self.transpose else feat def forward(self, x): if self.transpose: @@ -100,33 +108,50 @@ class GFSQ(nn.Module): embed_onehot_tmp = F.one_hot(ind.long(), self.n_ind) embed_onehot = embed_onehot_tmp.to(x.dtype) del embed_onehot_tmp - e_mean = torch.mean(embed_onehot, dim=[0,1]) + e_mean = torch.mean(embed_onehot, dim=[0, 1]) # e_mean = e_mean / (e_mean.sum(dim=1) + self.eps).unsqueeze(1) torch.div(e_mean, (e_mean.sum(dim=1) + self.eps).unsqueeze(1), out=e_mean) 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, + feat.transpose_(1, 2) if self.transpose else feat, perplexity, None, - ind.transpose_(1,2) if self.transpose else ind, + ind.transpose_(1, 2) if self.transpose else ind, ) + class DVAEDecoder(nn.Module): - def __init__(self, idim: int, odim: int, - n_layer = 12, bn_dim = 64, hidden = 256, - kernel = 7, dilation = 2, up = False - ): + def __init__( + self, + idim: int, + odim: int, + 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) + 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.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: torch.Tensor, conditioning=None) -> torch.Tensor: @@ -144,14 +169,20 @@ class DVAEDecoder(nn.Module): class DVAE(nn.Module): def __init__( - self, decoder_config, vq_config, dim=512, coef: Optional[str] = None, + self, + decoder_config, + vq_config, + dim=512, + coef: Optional[str] = None, ): super().__init__() if coef is None: coef = torch.rand(100) else: - coef = torch.from_numpy(np.copy(np.frombuffer(b14.decode_from_string(coef), dtype=np.float32))) - self.register_buffer('coef', coef.unsqueeze(0).unsqueeze_(2)) + coef = torch.from_numpy( + np.copy(np.frombuffer(b14.decode_from_string(coef), dtype=np.float32)) + ) + self.register_buffer("coef", coef.unsqueeze(0).unsqueeze_(2)) self.decoder = DVAEDecoder(**decoder_config) self.out_conv = nn.Conv1d(dim, 100, 3, 1, 1, bias=False) @@ -159,9 +190,11 @@ class DVAE(nn.Module): self.vq_layer = GFSQ(**vq_config) else: self.vq_layer = None - + def __repr__(self) -> str: - return b14.encode_to_string(self.coef.cpu().numpy().astype(np.float32).tobytes()) + return b14.encode_to_string( + self.coef.cpu().numpy().astype(np.float32).tobytes() + ) def forward(self, inp: torch.Tensor) -> torch.Tensor: with torch.no_grad(): @@ -171,9 +204,13 @@ class DVAE(nn.Module): else: vq_feats = inp.detach().clone() - vq_feats = vq_feats.view( - (vq_feats.size(0), 2, vq_feats.size(1)//2, vq_feats.size(2)), - ).permute(0, 2, 3, 1).flatten(2) + vq_feats = ( + vq_feats.view( + (vq_feats.size(0), 2, vq_feats.size(1) // 2, vq_feats.size(2)), + ) + .permute(0, 2, 3, 1) + .flatten(2) + ) dec_out = self.out_conv( self.decoder( diff --git a/ChatTTS/model/gpt.py b/ChatTTS/model/gpt.py index a9427b4..0f1241e 100644 --- a/ChatTTS/model/gpt.py +++ b/ChatTTS/model/gpt.py @@ -1,4 +1,5 @@ import os + os.environ["TOKENIZERS_PARALLELISM"] = "false" """ https://stackoverflow.com/questions/62691279/how-to-disable-tokenizers-parallelism-true-false-warning @@ -40,13 +41,13 @@ from ..utils import del_all class GPT(nn.Module): def __init__( - self, - gpt_config: omegaconf.DictConfig, + self, + gpt_config: omegaconf.DictConfig, num_audio_tokens: int, num_text_tokens: int, num_vq=4, device=torch.device("cpu"), - logger=logging.getLogger(__name__) + logger=logging.getLogger(__name__), ): super().__init__() @@ -58,49 +59,68 @@ class GPT(nn.Module): self.num_vq = num_vq self.num_audio_tokens = num_audio_tokens - self.gpt = self._build_llama(gpt_config, self.device_gpt) + self.gpt = self._build_llama(gpt_config, self.device_gpt) self.model_dim = int(self.gpt.config.hidden_size) self.emb_code = nn.ModuleList( - [nn.Embedding( - num_audio_tokens, self.model_dim, device=self.device_gpt, - ) for _ in range(num_vq)], + [ + nn.Embedding( + num_audio_tokens, + self.model_dim, + device=self.device_gpt, + ) + for _ in range(num_vq) + ], + ) + self.emb_text = nn.Embedding( + num_text_tokens, self.model_dim, device=self.device_gpt ) - self.emb_text = nn.Embedding(num_text_tokens, self.model_dim, device=self.device_gpt) self.head_text = weight_norm( nn.Linear( - self.model_dim, num_text_tokens, bias=False, device=device, + self.model_dim, + num_text_tokens, + bias=False, + device=device, ), - name='weight', + name="weight", ) self.head_code = nn.ModuleList( - [weight_norm( - nn.Linear( - self.model_dim, num_audio_tokens, bias=False, device=device, - ), - name='weight', - ) for _ in range(self.num_vq)], + [ + weight_norm( + nn.Linear( + self.model_dim, + num_audio_tokens, + bias=False, + device=device, + ), + name="weight", + ) + for _ in range(self.num_vq) + ], ) - + class Context: def __init__(self): self._interrupt = False def set(self, v: bool): self._interrupt = v - + def get(self) -> bool: return self._interrupt - - def _build_llama(self, config: omegaconf.DictConfig, device: torch.device) -> LlamaModel: + def _build_llama( + self, config: omegaconf.DictConfig, device: torch.device + ) -> LlamaModel: model = LlamaModel(LlamaConfig(**config)) del model.embed_tokens return model.to(device) - - def __call__(self, input_ids: torch.Tensor, text_mask: torch.Tensor) -> torch.Tensor: + + def __call__( + self, input_ids: torch.Tensor, text_mask: torch.Tensor + ) -> torch.Tensor: """ get_emb """ @@ -111,65 +131,89 @@ class GPT(nn.Module): get_emb """ - emb_text: torch.Tensor = self.emb_text(input_ids[text_mask].narrow(1, 0, 1).squeeze_(1).to(self.device_gpt)) + emb_text: torch.Tensor = self.emb_text( + input_ids[text_mask].narrow(1, 0, 1).squeeze_(1).to(self.device_gpt) + ) text_mask_inv = text_mask.logical_not().to(self.device_gpt) masked_input_ids: torch.Tensor = input_ids[text_mask_inv].to(self.device_gpt) - emb_code = [self.emb_code[i](masked_input_ids[:, i]) for i in range(self.num_vq)] + emb_code = [ + self.emb_code[i](masked_input_ids[:, 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 = 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_inv] = emb_code.to(emb.dtype) del emb_text, emb_code, text_mask_inv return emb - + @dataclass(repr=False, eq=False) - class _GenerationInputs(): + class _GenerationInputs: position_ids: torch.Tensor cache_position: torch.Tensor use_cache: bool - input_ids: Optional[torch.Tensor]=None - past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]]=None - attention_mask: Optional[torch.Tensor]=None - inputs_embeds: Optional[torch.Tensor]=None + input_ids: Optional[torch.Tensor] = None + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + attention_mask: Optional[torch.Tensor] = None + inputs_embeds: Optional[torch.Tensor] = None def to(self, device: torch.device): - if self.attention_mask is not None: self.attention_mask = self.attention_mask.to(device) - if self.position_ids is not None: self.position_ids = self.position_ids.to(device) - if self.inputs_embeds is not None: self.inputs_embeds = self.inputs_embeds.to(device) - if self.cache_position is not None: self.cache_position = self.cache_position.to(device) + if self.attention_mask is not None: + self.attention_mask = self.attention_mask.to(device) + if self.position_ids is not None: + self.position_ids = self.position_ids.to(device) + if self.inputs_embeds is not None: + self.inputs_embeds = self.inputs_embeds.to(device) + if self.cache_position is not None: + self.cache_position = self.cache_position.to(device) def _prepare_generation_inputs( self, input_ids: torch.Tensor, - past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]]=None, - attention_mask: Optional[torch.Tensor]=None, - inputs_embeds: Optional[torch.Tensor]=None, - cache_position: Optional[torch.Tensor]=None, - position_ids: Optional[torch.Tensor]=None, - use_cache = True, + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + attention_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + cache_position: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + use_cache=True, ) -> _GenerationInputs: # 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) + 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() + 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) + 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) + 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] @@ -179,7 +223,10 @@ class GPT(nn.Module): # 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]: + 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. @@ -202,9 +249,13 @@ class GPT(nn.Module): if past_key_values: position_ids = position_ids[:, -input_ids.shape[1] :] - input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1] + 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) + cache_position = torch.arange( + past_length, past_length + input_length, device=input_ids.device + ) else: cache_position = cache_position[-input_length:] @@ -232,7 +283,7 @@ class GPT(nn.Module): return model_inputs @dataclass(repr=False, eq=False) - class GenerationOutputs(): + class GenerationOutputs: ids: List[torch.Tensor] attentions: List[Optional[Tuple[torch.FloatTensor, ...]]] hiddens: List[torch.Tensor] @@ -242,7 +293,6 @@ class GPT(nn.Module): del_all(self.attentions) del_all(self.hiddens) - def _prepare_generation_outputs( self, inputs_ids: torch.Tensor, @@ -252,13 +302,17 @@ class GPT(nn.Module): hiddens: List[torch.Tensor], infer_text: bool, ) -> GenerationOutputs: - inputs_ids = [inputs_ids[idx].narrow(0, start_idx, i) for idx, i in enumerate(end_idx)] + inputs_ids = [ + inputs_ids[idx].narrow(0, start_idx, i) for idx, i in enumerate(end_idx) + ] if infer_text: inputs_ids = [i.narrow(1, 0, 1).squeeze_(1) for i in inputs_ids] if len(hiddens) > 0: hiddens = torch.stack(hiddens, 1) - hiddens = [hiddens[idx].narrow(0, 0, i) for idx, i in enumerate(end_idx.int())] + hiddens = [ + hiddens[idx].narrow(0, 0, i) for idx, i in enumerate(end_idx.int()) + ] return self.GenerationOutputs( ids=inputs_ids, @@ -266,16 +320,15 @@ class GPT(nn.Module): hiddens=hiddens, ) - def generate( - self, - emb: torch.Tensor, - inputs_ids: torch.Tensor, - temperature: torch.Tensor, - eos_token: Union[int, torch.Tensor], - attention_mask = None, - max_new_token = 2048, - min_new_token = 0, + self, + emb: torch.Tensor, + inputs_ids: torch.Tensor, + temperature: torch.Tensor, + eos_token: Union[int, torch.Tensor], + attention_mask=None, + max_new_token=2048, + min_new_token=0, logits_warpers: List[LogitsWarper] = [], logits_processors: List[CustomRepetitionPenaltyLogitsProcessorRepeat] = [], infer_text=False, @@ -284,35 +337,49 @@ class GPT(nn.Module): stream=False, context=Context(), ): - + with torch.no_grad(): attentions: List[Optional[Tuple[torch.FloatTensor, ...]]] = [] hiddens = [] - start_idx, end_idx = inputs_ids.shape[1], torch.zeros(inputs_ids.shape[0], device=inputs_ids.device, dtype=torch.long) + 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.unsqueeze_(0).expand(inputs_ids.shape[0], -1).contiguous().view(-1, 1) + + temperature = ( + temperature.unsqueeze_(0) + .expand(inputs_ids.shape[0], -1) + .contiguous() + .view(-1, 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) + 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 + attention_mask_cache[:, : attention_mask.shape[1]] = attention_mask with tqdm( total=max_new_token, desc="text" if infer_text else "code", - bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt}(max) [{elapsed}, {rate_fmt}{postfix}]', + bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}(max) [{elapsed}, {rate_fmt}{postfix}]", ) as pbar: past_key_values = None for i in range(max_new_token): model_input = self._prepare_generation_inputs( - inputs_ids, - past_key_values, - attention_mask_cache[:, :inputs_ids.shape[1]], + inputs_ids, + past_key_values, + attention_mask_cache[:, : inputs_ids.shape[1]], use_cache=True, ) @@ -320,9 +387,12 @@ class GPT(nn.Module): del emb inputs_ids_emb = model_input.input_ids.to(self.device_gpt) if infer_text: - emb: torch.Tensor = self.emb_text(inputs_ids_emb[:,:,0]) + emb: torch.Tensor = self.emb_text(inputs_ids_emb[:, :, 0]) else: - code_emb = [self.emb_code[i](inputs_ids_emb[:,:,i]) for i in range(self.num_vq)] + code_emb = [ + self.emb_code[i](inputs_ids_emb[:, :, i]) + for i in range(self.num_vq) + ] emb = torch.stack(code_emb, 3).sum(3) del inputs_ids_emb, model_input.input_ids model_input.inputs_embeds = emb @@ -340,7 +410,7 @@ class GPT(nn.Module): ) del_all(model_input) attentions.append(outputs.attentions) - hidden_states = outputs.last_hidden_state.to(self.device) # 🐻 + hidden_states = outputs.last_hidden_state.to(self.device) # 🐻 past_key_values = outputs.past_key_values del_all(outputs) if return_hidden: @@ -352,9 +422,12 @@ class GPT(nn.Module): else: # logits = torch.stack([self.head_code[i](hidden_states) for i in range(self.num_vq)], 3) logits = torch.empty( - hidden_states.size(0), hidden_states.size(1), - self.num_audio_tokens, self.num_vq, - dtype=torch.float, device=self.device, + hidden_states.size(0), + hidden_states.size(1), + self.num_audio_tokens, + self.num_vq, + dtype=torch.float, + device=self.device, ) for i in range(self.num_vq): x: torch.Tensor = self.head_code[i](hidden_states) @@ -371,7 +444,8 @@ class GPT(nn.Module): # logits_token = rearrange(inputs_ids[:, start_idx:], "b c n -> (b n) c") inputs_ids_sliced = inputs_ids[:, start_idx:].permute(0, 2, 1) logits_token = inputs_ids_sliced.reshape( - inputs_ids_sliced.size(0)*inputs_ids_sliced.size(1), -1, + inputs_ids_sliced.size(0) * inputs_ids_sliced.size(1), + -1, ).to(self.device) else: logits_token = inputs_ids[:, start_idx:, 0].to(self.device) @@ -393,7 +467,9 @@ class GPT(nn.Module): del logits - idx_next = torch.multinomial(scores, num_samples=1).to(finish.device) + idx_next = torch.multinomial(scores, num_samples=1).to( + finish.device + ) if not infer_text: # idx_next = rearrange(idx_next, "(b n) 1 -> b n", n=self.num_vq) @@ -401,12 +477,20 @@ class GPT(nn.Module): finish_or = (idx_next == eos_token).any(1) finish.logical_or_(finish_or) del finish_or - inputs_ids_tmp = torch.cat([inputs_ids, idx_next.unsqueeze_(1)], 1) + inputs_ids_tmp = torch.cat( + [inputs_ids, idx_next.unsqueeze_(1)], 1 + ) else: finish_or = (idx_next == eos_token).any(1) finish.logical_or_(finish_or) del finish_or - inputs_ids_tmp = torch.cat([inputs_ids, idx_next.unsqueeze_(-1).expand(-1, -1, self.num_vq)], 1) + inputs_ids_tmp = torch.cat( + [ + inputs_ids, + idx_next.unsqueeze_(-1).expand(-1, -1, self.num_vq), + ], + 1, + ) del inputs_ids inputs_ids = inputs_ids_tmp @@ -416,27 +500,42 @@ class GPT(nn.Module): minus_prev_end_index = end_idx.neg() end_idx.add_((finish.logical_not().to(end_idx.device)).int()) if stream: - if end_idx.all() and (end_idx%24 == 0).any() and minus_prev_end_index.add_(end_idx).any(): + if ( + end_idx.all() + and (end_idx % 24 == 0).any() + and minus_prev_end_index.add_(end_idx).any() + ): self.logger.debug("yield stream result, end: %d", end_idx) yield self._prepare_generation_outputs( - inputs_ids, start_idx, end_idx, attentions, hiddens, + inputs_ids, + start_idx, + end_idx, + attentions, + hiddens, infer_text, ) del minus_prev_end_index - if finish.all() or context.get(): break + if finish.all() or context.get(): + break pbar.update(1) if not finish.all(): if context.get(): - self.logger.warning('generation is interrupted') + self.logger.warning("generation is interrupted") else: - self.logger.warning(f'incomplete result. hit max_new_token: {max_new_token}') + self.logger.warning( + f"incomplete result. hit max_new_token: {max_new_token}" + ) del finish yield self._prepare_generation_outputs( - inputs_ids, start_idx, end_idx, attentions, hiddens, + inputs_ids, + start_idx, + end_idx, + attentions, + hiddens, infer_text, ) diff --git a/ChatTTS/model/processors.py b/ChatTTS/model/processors.py index 589c679..02787e9 100644 --- a/ChatTTS/model/processors.py +++ b/ChatTTS/model/processors.py @@ -1,28 +1,33 @@ import torch import torch.nn.functional as F - -class CustomRepetitionPenaltyLogitsProcessorRepeat(): + +class CustomRepetitionPenaltyLogitsProcessorRepeat: def __init__(self, penalty: float, max_input_ids: int, past_window: int): if not isinstance(penalty, float) or not (penalty > 0): - raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}") + 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:] + 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 + freq[self.max_input_ids :] = 0 alpha = self.penalty**freq scores = scores.contiguous() - scores = torch.where(scores < 0, scores*alpha, scores/alpha) + scores = torch.where(scores < 0, scores * alpha, scores / alpha) return scores - + + """class CustomRepetitionPenaltyLogitsProcessor(): def __init__(self, penalty: float, max_input_ids: int, past_window: int): diff --git a/ChatTTS/norm.py b/ChatTTS/norm.py index f4c87ef..bf6e7b8 100644 --- a/ChatTTS/norm.py +++ b/ChatTTS/norm.py @@ -17,8 +17,11 @@ def _find_index(table: np.ndarray, val: np.uint16): return i return -1 + @jit -def _fast_replace(table: np.ndarray, text: bytes) -> Tuple[np.ndarray, List[Tuple[str, str]]]: +def _fast_replace( + table: np.ndarray, text: bytes +) -> Tuple[np.ndarray, List[Tuple[str, str]]]: result = np.frombuffer(text, dtype=np.uint16).copy() replaced_words = [] for i in range(result.size): @@ -30,6 +33,7 @@ def _fast_replace(table: np.ndarray, text: bytes) -> Tuple[np.ndarray, List[Tupl replaced_words.append((chr(ch), chr(repl_char))) return result, replaced_words + class Normalizer: def __init__(self, map_file_path: str, logger=logging.getLogger(__name__)): self.logger = logger @@ -53,75 +57,79 @@ class Normalizer: """ self.coding = "utf-16-le" if sys.byteorder == "little" else "utf-16-be" - self.accept_pattern = re.compile(r'[^\u4e00-\u9fffA-Za-z，。、,\. ]') - self.sub_pattern = re.compile(r'\[uv_break\]|\[laugh\]|\[lbreak\]') - self.chinese_char_pattern = re.compile(r'[\u4e00-\u9fff]') - self.english_word_pattern = re.compile(r'\b[A-Za-z]+\b') - self.character_simplifier = str.maketrans({ - '：': '，', - '；': '，', - '！': '。', - '（': '，', - '）': '，', - '【': '，', - '】': '，', - '『': '，', - '』': '，', - '「': '，', - '」': '，', - '《': '，', - '》': '，', - '－': '，', - '‘': '', - '“': '', - '’': '', - '”': '', - ':': ',', - ';': ',', - '!': '.', - '(': ',', - ')': ',', - '[': ',', - ']': ',', - '>': ',', - '<': ',', - '-': ',', - }) - self.halfwidth_2_fullwidth = str.maketrans({ - '!': '！', - '"': '“', - "'": '‘', - '#': '＃', - '$': '＄', - '%': '％', - '&': '＆', - '(': '（', - ')': '）', - ',': '，', - '-': '－', - '*': '＊', - '+': '＋', - '.': '。', - '/': '／', - ':': '：', - ';': '；', - '<': '＜', - '=': '＝', - '>': '＞', - '?': '？', - '@': '＠', + self.accept_pattern = re.compile(r"[^\u4e00-\u9fffA-Za-z，。、,\. ]") + self.sub_pattern = re.compile(r"\[uv_break\]|\[laugh\]|\[lbreak\]") + self.chinese_char_pattern = re.compile(r"[\u4e00-\u9fff]") + self.english_word_pattern = re.compile(r"\b[A-Za-z]+\b") + self.character_simplifier = str.maketrans( + { + "：": "，", + "；": "，", + "！": "。", + "（": "，", + "）": "，", + "【": "，", + "】": "，", + "『": "，", + "』": "，", + "「": "，", + "」": "，", + "《": "，", + "》": "，", + "－": "，", + "‘": "", + "“": "", + "’": "", + "”": "", + ":": ",", + ";": ",", + "!": ".", + "(": ",", + ")": ",", + "[": ",", + "]": ",", + ">": ",", + "<": ",", + "-": ",", + } + ) + self.halfwidth_2_fullwidth = str.maketrans( + { + "!": "！", + '"': "“", + "'": "‘", + "#": "＃", + "$": "＄", + "%": "％", + "&": "＆", + "(": "（", + ")": "）", + ",": "，", + "-": "－", + "*": "＊", + "+": "＋", + ".": "。", + "/": "／", + ":": "：", + ";": "；", + "<": "＜", + "=": "＝", + ">": "＞", + "?": "？", + "@": "＠", # '[': '［', - '\\': '＼', + "\\": "＼", # ']': '］', - '^': '＾', + "^": "＾", # '_': '＿', - '`': '｀', - '{': '｛', - '|': '｜', - '}': '｝', - '~': '～' - }) - + "`": "｀", + "{": "｛", + "|": "｜", + "}": "｝", + "~": "～", + } + ) + def __call__( self, text: str, @@ -133,11 +141,11 @@ class Normalizer: _lang = self._detect_language(text) if lang is None else lang if _lang in self.normalizers: text = self.normalizers[_lang](text) - if _lang == 'zh': + if _lang == "zh": text = self._apply_half2full_map(text) invalid_characters = self._count_invalid_characters(text) if len(invalid_characters): - self.logger.warning(f'found invalid characters: {invalid_characters}') + self.logger.warning(f"found invalid characters: {invalid_characters}") text = self._apply_character_map(text) if do_homophone_replacement: arr, replaced_words = _fast_replace( @@ -146,11 +154,10 @@ class Normalizer: ) if replaced_words: text = arr.tobytes().decode(self.coding) - repl_res = ', '.join([f'{_[0]}->{_[1]}' for _ in replaced_words]) - self.logger.info(f'replace homophones: {repl_res}') + repl_res = ", ".join([f"{_[0]}->{_[1]}" for _ in replaced_words]) + self.logger.info(f"replace homophones: {repl_res}") return text - def register(self, name: str, normalizer: Callable[[str], str]) -> bool: if name in self.normalizers: self.logger.warning(f"name {name} has been registered") @@ -169,25 +176,25 @@ class Normalizer: def unregister(self, name: str): if name in self.normalizers: del self.normalizers[name] - + def destroy(self): del_all(self.normalizers) del self.homophones_map def _load_homophones_map(self, map_file_path: str) -> np.ndarray: - with open(map_file_path, 'r', encoding='utf-8') as f: + with open(map_file_path, "r", encoding="utf-8") as f: homophones_map: Dict[str, str] = json.load(f) map = np.empty((2, len(homophones_map)), dtype=np.uint32) for i, k in enumerate(homophones_map.keys()): map[:, i] = (ord(k), ord(homophones_map[k])) del homophones_map return map - + def _count_invalid_characters(self, s: str): - s = self.sub_pattern.sub('', s) + s = self.sub_pattern.sub("", s) non_alphabetic_chinese_chars = self.accept_pattern.findall(s) return set(non_alphabetic_chinese_chars) - + def _apply_half2full_map(self, text: str) -> str: return text.translate(self.halfwidth_2_fullwidth) diff --git a/ChatTTS/utils/dl.py b/ChatTTS/utils/dl.py index 2044015..1bd66da 100644 --- a/ChatTTS/utils/dl.py +++ b/ChatTTS/utils/dl.py @@ -7,6 +7,7 @@ from mmap import mmap, ACCESS_READ from .log import logger + def sha256(fileno: int) -> str: data = mmap(fileno, 0, access=ACCESS_READ) h = hashlib.sha256(data).hexdigest() diff --git a/ChatTTS/utils/gpu.py b/ChatTTS/utils/gpu.py index 17a3067..550c6cc 100644 --- a/ChatTTS/utils/gpu.py +++ b/ChatTTS/utils/gpu.py @@ -1,8 +1,8 @@ - import torch from .log import logger + def select_device(min_memory=2047): if torch.cuda.is_available(): available_gpus = [] @@ -11,17 +11,19 @@ def select_device(min_memory=2047): 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}') + device = torch.device(f"cuda:{selected_gpu}") free_memory_mb = max_free_memory / (1024 * 1024) if free_memory_mb < min_memory: - logger.get_logger().warning(f'GPU {selected_gpu} has {round(free_memory_mb, 2)} MB memory left. Switching to CPU.') - device = torch.device('cpu') + logger.get_logger().warning( + f"GPU {selected_gpu} has {round(free_memory_mb, 2)} MB memory left. Switching to CPU." + ) + device = torch.device("cpu") elif torch.backends.mps.is_available(): # For Apple M1/M2 chips with Metal Performance Shaders - logger.get_logger().info('apple GPU found, using MPS.') - device = torch.device('mps') + logger.get_logger().info("apple GPU found, using MPS.") + device = torch.device("mps") else: - logger.get_logger().warning('no GPU found, use CPU instead') - device = torch.device('cpu') + logger.get_logger().warning("no GPU found, use CPU instead") + device = torch.device("cpu") return device diff --git a/ChatTTS/utils/io.py b/ChatTTS/utils/io.py index b787d21..b37f939 100644 --- a/ChatTTS/utils/io.py +++ b/ChatTTS/utils/io.py @@ -1,4 +1,3 @@ - import os import logging from typing import Union @@ -6,16 +5,20 @@ from dataclasses import is_dataclass from .log import logger + def get_latest_modified_file(directory): - files = [os.path.join(directory, f) for f in os.listdir(directory)] + files = [os.path.join(directory, f) for f in os.listdir(directory)] if not files: - logger.get_logger().log(logging.WARNING, f'no files found in the directory: {directory}') + logger.get_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 + def del_all(d: Union[dict, list]): if is_dataclass(d): for k in list(vars(d).keys()): @@ -39,4 +42,3 @@ def del_all(d: Union[dict, list]): del x else: del d - diff --git a/ChatTTS/utils/log.py b/ChatTTS/utils/log.py index 382d79f..1fd9b93 100644 --- a/ChatTTS/utils/log.py +++ b/ChatTTS/utils/log.py @@ -1,14 +1,16 @@ import logging from pathlib import Path -class Logger(): + +class Logger: def __init__(self, logger=logging.getLogger(Path(__file__).parent.name)): self.logger = logger def set_logger(self, logger: logging.Logger): self.logger = logger - + def get_logger(self) -> logging.Logger: return self.logger + logger = Logger() diff --git a/examples/cmd/run.py b/examples/cmd/run.py index e910a0a..12fb8f6 100644 --- a/examples/cmd/run.py +++ b/examples/cmd/run.py @@ -16,6 +16,7 @@ from tools.logger import get_logger logger = get_logger("Command") + def save_wav_file(wav, index): wav_filename = f"output_audio_{index}.wav" with wave.open(wav_filename, "wb") as wf: @@ -25,6 +26,7 @@ def save_wav_file(wav, index): wf.writeframes(unsafe_float_to_int16(wav)) logger.info(f"Audio saved to {wav_filename}") + def main(texts: list[str]): logger.info("Text input: %s", str(texts)) @@ -42,10 +44,15 @@ def main(texts: list[str]): for index, wav in enumerate(wavs): save_wav_file(wav, index) + if __name__ == "__main__": logger.info("Starting the TTS application...") - parser = argparse.ArgumentParser(description='ChatTTS Command', usage="--stream hello, my name is bob.") - parser.add_argument("text", help="Original text", default='YOUR TEXT HERE', nargs='*') + parser = argparse.ArgumentParser( + description="ChatTTS Command", usage="--stream hello, my name is bob." + ) + parser.add_argument( + "text", help="Original text", default="YOUR TEXT HERE", nargs="*" + ) args = parser.parse_args() main(args.text) logger.info("TTS application finished.") diff --git a/examples/ipynb/colab.ipynb b/examples/ipynb/colab.ipynb index bd6f509..5a3900d 100644 --- a/examples/ipynb/colab.ipynb +++ b/examples/ipynb/colab.ipynb @@ -1,474 +1,484 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "id": "xYJFXKP9xhQM" - }, - "source": [ - "## Clone Repo" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "hegwDOfffwzw" - }, - "outputs": [], - "source": [ - "!cd /content\n", - "!rm -rf sample_data ChatTTS\n", - "!git clone https://github.com/2noise/ChatTTS.git\n", - "!pip install -r /content/ChatTTS/requirements.txt\n", - "!ldconfig /usr/lib64-nvidia" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "zdzEFoknxqTH" - }, - "source": [ - "## Import Libs" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "lDSQ6Xf-bSre" - }, - "outputs": [], - "source": [ - "import torch\n", - "torch._dynamo.config.cache_size_limit = 64\n", - "torch._dynamo.config.suppress_errors = True\n", - "torch.set_float32_matmul_precision('high')\n", - "\n", - "from ChatTTS import ChatTTS\n", - "from ChatTTS.tools.logger import get_logger\n", - "from ChatTTS.tools.normalizer import normalizer_en_nemo_text, normalizer_zh_tn\n", - "from IPython.display import Audio" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "vBzG5gxcbSrf" - }, - "source": [ - "## Load Models" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "e0QSkngRbSrg" - }, - "outputs": [], - "source": [ - "logger = get_logger(\"ChatTTS\", format_root=True)\n", - "chat = ChatTTS.Chat(logger)\n", - "\n", - "# try to load normalizer\n", - "try:\n", - " chat.normalizer.register(\"en\", normalizer_en_nemo_text())\n", - "except ValueError as e:\n", - " logger.error(e)\n", - "except:\n", - " logger.warning('Package nemo_text_processing not found!')\n", - " logger.warning(\n", - " 'Run: conda install -c conda-forge pynini=2.1.5 && pip install nemo_text_processing',\n", - " )\n", - "try:\n", - " chat.normalizer.register(\"zh\", normalizer_zh_tn())\n", - "except ValueError as e:\n", - " logger.error(e)\n", - "except:\n", - " logger.warning('Package WeTextProcessing not found!')\n", - " logger.warning(\n", - " 'Run: conda install -c conda-forge pynini=2.1.5 && pip install WeTextProcessing',\n", - " )" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "3Ty427FZNH30" - }, - "source": [ - "### Here are three choices for loading models:" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "NInF7Lk1NH30" - }, - "source": [ - "#### 1. Load models from Hugging Face:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "VVtNlNosNH30" - }, - "outputs": [], - "source": [ - "# use force_redownload=True if the weights have been updated.\n", - "chat.load(source='huggingface', force_redownload=True)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "AhBD5WUPNH30" - }, - "source": [ - "#### 2. Load models from local directories 'asset' and 'config':" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "83UwV6SGNH31" - }, - "outputs": [], - "source": [ - "chat.load()\n", - "# chat.load(source='local') same as above" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "c0qjGPNkNH31" - }, - "source": [ - "#### 3. Load models from a custom path:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "oCSBx0Q7NH31" - }, - "outputs": [], - "source": [ - "# write the model path into custom_path\n", - "chat.load(source='custom', custom_path='YOUR CUSTOM PATH')" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "VoEki3XMNH31" - }, - "source": [ - "### You can also unload models to save the memory" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "3FdsTSxoNH31" - }, - "outputs": [], - "source": [ - "chat.unload()" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "bAUs0rGQbSrh" - }, - "source": [ - "## Inference" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "NPZ2SFksbSrh" - }, - "source": [ - "### Batch infer" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "Su9FmUYAbSrh" - }, - "outputs": [], - "source": [ - "texts = [\"So we found being competitive and collaborative was a huge way of staying motivated towards our goals, so one person to call when you fall off, one person who gets you back on then one person to actually do the activity with.\",]*3 \\\n", - " + [\"我觉得像我们这些写程序的人，他，我觉得多多少少可能会对开源有一种情怀在吧我觉得开源是一个很好的形式。现在其实最先进的技术掌握在一些公司的手里的话，就他们并不会轻易的开放给所有的人用。\"]*3\n", - "\n", - "wavs = chat.infer(texts)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "YQRwB8lpbSri" - }, - "outputs": [], - "source": [ - "Audio(wavs[0], rate=24_000, autoplay=True)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "LuFG6m7AbSri" - }, - "outputs": [], - "source": [ - "Audio(wavs[3], rate=24_000, autoplay=True)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "oLhAGvkfbSrj" - }, - "source": [ - "### Custom params" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "kma0HBEBbSrj" - }, - "outputs": [], - "source": [ - "params_infer_code = ChatTTS.Chat.InferCodeParams(\n", - " prompt='[speed_5]',\n", - " temperature=.3,\n", - ")\n", - "params_refine_text = ChatTTS.Chat.RefineTextParams(\n", - " prompt='[oral_2][laugh_0][break_6]',\n", - ")\n", - "\n", - "wav = chat.infer('四川美食可多了，有麻辣火锅、宫保鸡丁、麻婆豆腐、担担面、回锅肉、夫妻肺片等，每样都让人垂涎三尺。', \\\n", - " params_refine_text=params_refine_text, params_infer_code=params_infer_code)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "Nl_mT9KpbSrj" - }, - "outputs": [], - "source": [ - "Audio(wav[0], rate=24_000, autoplay=True)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "JfAba-tTbSrk" - }, - "source": [ - "### fix random speaker" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "Qh7dcWrAbSrk" - }, - "outputs": [], - "source": [ - "rand_spk = chat.sample_random_speaker()\n", - "params_infer_code = ChatTTS.Chat.InferCodeParams(\n", - " spk_emb=rand_spk,\n", - ")\n", - "\n", - "wav = chat.infer('四川美食确实以辣闻名，但也有不辣的选择。比如甜水面、赖汤圆、蛋烘糕、叶儿粑等，这些小吃口味温和，甜而不腻，也很受欢迎。', \\\n", - " params_refine_text=params_refine_text, params_infer_code=params_infer_code)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "0ljWDWzabSrk" - }, - "outputs": [], - "source": [ - "Audio(wav[0], rate=24_000, autoplay=True)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "u1q-BcUKbSrl" - }, - "source": [ - "### Two stage control" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "3hAAc0lJbSrl" - }, - "outputs": [], - "source": [ - "text = \"So we found being competitive and collaborative was a huge way of staying motivated towards our goals, so one person to call when you fall off, one person who gets you back on then one person to actually do the activity with.\"\n", - "refined_text = chat.infer(text, refine_text_only=True)\n", - "refined_text" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "0GVJxhd3BKQX" - }, - "outputs": [], - "source": [ - "wav = chat.infer(refined_text, skip_refine_text=True)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "ngyMht74BicY" - }, - "outputs": [], - "source": [ - "Audio(wav[0], rate=24_000, autoplay=True)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "GG5AMbQbbSrl" - }, - "source": [ - "## LLM Call" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "3rkfwc3UbSrl" - }, - "outputs": [], - "source": [ - "from ChatTTS.tools.llm import ChatOpenAI\n", - "\n", - "API_KEY = ''\n", - "client = ChatOpenAI(api_key=API_KEY,\n", - " base_url=\"https://api.deepseek.com\",\n", - " model=\"deepseek-chat\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "TTkIsXozbSrm" - }, - "outputs": [], - "source": [ - "user_question = '四川有哪些好吃的美食呢?'" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "3yT8uNz-RVy1" - }, - "outputs": [], - "source": [ - "text = client.call(user_question, prompt_version = 'deepseek')\n", - "text" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "6qddpv7lRW-3" - }, - "outputs": [], - "source": [ - "text = client.call(text, prompt_version = 'deepseek_TN')\n", - "text" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "qNhCJG4VbSrm" - }, - "outputs": [], - "source": [ - "wav = chat.infer(text)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "Wq1XQHmFRQI3" - }, - "outputs": [], - "source": [ - "Audio(wav[0], rate=24_000, autoplay=True)" - ] - } - ], - "metadata": { - "accelerator": "GPU", - "colab": { - "collapsed_sections": [ - "bAUs0rGQbSrh" - ], - "gpuType": "T4", - "provenance": [] - }, - "kernelspec": { - "display_name": "Python 3", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.10.8" - } + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "xYJFXKP9xhQM" + }, + "source": [ + "## Clone Repo" + ] }, - "nbformat": 4, - "nbformat_minor": 0 + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "hegwDOfffwzw" + }, + "outputs": [], + "source": [ + "!cd /content\n", + "!rm -rf sample_data ChatTTS\n", + "!git clone https://github.com/2noise/ChatTTS.git\n", + "!pip install -r /content/ChatTTS/requirements.txt\n", + "!ldconfig /usr/lib64-nvidia" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "zdzEFoknxqTH" + }, + "source": [ + "## Import Libs" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "lDSQ6Xf-bSre" + }, + "outputs": [], + "source": [ + "import torch\n", + "\n", + "torch._dynamo.config.cache_size_limit = 64\n", + "torch._dynamo.config.suppress_errors = True\n", + "torch.set_float32_matmul_precision(\"high\")\n", + "\n", + "from ChatTTS import ChatTTS\n", + "from ChatTTS.tools.logger import get_logger\n", + "from ChatTTS.tools.normalizer import normalizer_en_nemo_text, normalizer_zh_tn\n", + "from IPython.display import Audio" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "vBzG5gxcbSrf" + }, + "source": [ + "## Load Models" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "e0QSkngRbSrg" + }, + "outputs": [], + "source": [ + "logger = get_logger(\"ChatTTS\", format_root=True)\n", + "chat = ChatTTS.Chat(logger)\n", + "\n", + "# try to load normalizer\n", + "try:\n", + " chat.normalizer.register(\"en\", normalizer_en_nemo_text())\n", + "except ValueError as e:\n", + " logger.error(e)\n", + "except:\n", + " logger.warning(\"Package nemo_text_processing not found!\")\n", + " logger.warning(\n", + " \"Run: conda install -c conda-forge pynini=2.1.5 && pip install nemo_text_processing\",\n", + " )\n", + "try:\n", + " chat.normalizer.register(\"zh\", normalizer_zh_tn())\n", + "except ValueError as e:\n", + " logger.error(e)\n", + "except:\n", + " logger.warning(\"Package WeTextProcessing not found!\")\n", + " logger.warning(\n", + " \"Run: conda install -c conda-forge pynini=2.1.5 && pip install WeTextProcessing\",\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "3Ty427FZNH30" + }, + "source": [ + "### Here are three choices for loading models:" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "NInF7Lk1NH30" + }, + "source": [ + "#### 1. Load models from Hugging Face:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "VVtNlNosNH30" + }, + "outputs": [], + "source": [ + "# use force_redownload=True if the weights have been updated.\n", + "chat.load(source=\"huggingface\", force_redownload=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "AhBD5WUPNH30" + }, + "source": [ + "#### 2. Load models from local directories 'asset' and 'config':" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "83UwV6SGNH31" + }, + "outputs": [], + "source": [ + "chat.load()\n", + "# chat.load(source='local') same as above" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "c0qjGPNkNH31" + }, + "source": [ + "#### 3. Load models from a custom path:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "oCSBx0Q7NH31" + }, + "outputs": [], + "source": [ + "# write the model path into custom_path\n", + "chat.load(source=\"custom\", custom_path=\"YOUR CUSTOM PATH\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "VoEki3XMNH31" + }, + "source": [ + "### You can also unload models to save the memory" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "3FdsTSxoNH31" + }, + "outputs": [], + "source": [ + "chat.unload()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "bAUs0rGQbSrh" + }, + "source": [ + "## Inference" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "NPZ2SFksbSrh" + }, + "source": [ + "### Batch infer" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Su9FmUYAbSrh" + }, + "outputs": [], + "source": [ + "texts = [\n", + " \"So we found being competitive and collaborative was a huge way of staying motivated towards our goals, so one person to call when you fall off, one person who gets you back on then one person to actually do the activity with.\",\n", + "] * 3 + [\n", + " \"我觉得像我们这些写程序的人，他，我觉得多多少少可能会对开源有一种情怀在吧我觉得开源是一个很好的形式。现在其实最先进的技术掌握在一些公司的手里的话，就他们并不会轻易的开放给所有的人用。\"\n", + "] * 3\n", + "\n", + "wavs = chat.infer(texts)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "YQRwB8lpbSri" + }, + "outputs": [], + "source": [ + "Audio(wavs[0], rate=24_000, autoplay=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "LuFG6m7AbSri" + }, + "outputs": [], + "source": [ + "Audio(wavs[3], rate=24_000, autoplay=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "oLhAGvkfbSrj" + }, + "source": [ + "### Custom params" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "kma0HBEBbSrj" + }, + "outputs": [], + "source": [ + "params_infer_code = ChatTTS.Chat.InferCodeParams(\n", + " prompt=\"[speed_5]\",\n", + " temperature=0.3,\n", + ")\n", + "params_refine_text = ChatTTS.Chat.RefineTextParams(\n", + " prompt=\"[oral_2][laugh_0][break_6]\",\n", + ")\n", + "\n", + "wav = chat.infer(\n", + " \"四川美食可多了，有麻辣火锅、宫保鸡丁、麻婆豆腐、担担面、回锅肉、夫妻肺片等，每样都让人垂涎三尺。\",\n", + " params_refine_text=params_refine_text,\n", + " params_infer_code=params_infer_code,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Nl_mT9KpbSrj" + }, + "outputs": [], + "source": [ + "Audio(wav[0], rate=24_000, autoplay=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "JfAba-tTbSrk" + }, + "source": [ + "### fix random speaker" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Qh7dcWrAbSrk" + }, + "outputs": [], + "source": [ + "rand_spk = chat.sample_random_speaker()\n", + "params_infer_code = ChatTTS.Chat.InferCodeParams(\n", + " spk_emb=rand_spk,\n", + ")\n", + "\n", + "wav = chat.infer(\n", + " \"四川美食确实以辣闻名，但也有不辣的选择。比如甜水面、赖汤圆、蛋烘糕、叶儿粑等，这些小吃口味温和，甜而不腻，也很受欢迎。\",\n", + " params_refine_text=params_refine_text,\n", + " params_infer_code=params_infer_code,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "0ljWDWzabSrk" + }, + "outputs": [], + "source": [ + "Audio(wav[0], rate=24_000, autoplay=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "u1q-BcUKbSrl" + }, + "source": [ + "### Two stage control" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "3hAAc0lJbSrl" + }, + "outputs": [], + "source": [ + "text = \"So we found being competitive and collaborative was a huge way of staying motivated towards our goals, so one person to call when you fall off, one person who gets you back on then one person to actually do the activity with.\"\n", + "refined_text = chat.infer(text, refine_text_only=True)\n", + "refined_text" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "0GVJxhd3BKQX" + }, + "outputs": [], + "source": [ + "wav = chat.infer(refined_text, skip_refine_text=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ngyMht74BicY" + }, + "outputs": [], + "source": [ + "Audio(wav[0], rate=24_000, autoplay=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "GG5AMbQbbSrl" + }, + "source": [ + "## LLM Call" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "3rkfwc3UbSrl" + }, + "outputs": [], + "source": [ + "from ChatTTS.tools.llm import ChatOpenAI\n", + "\n", + "API_KEY = \"\"\n", + "client = ChatOpenAI(\n", + " api_key=API_KEY, base_url=\"https://api.deepseek.com\", model=\"deepseek-chat\"\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "TTkIsXozbSrm" + }, + "outputs": [], + "source": [ + "user_question = \"四川有哪些好吃的美食呢?\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "3yT8uNz-RVy1" + }, + "outputs": [], + "source": [ + "text = client.call(user_question, prompt_version=\"deepseek\")\n", + "text" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "6qddpv7lRW-3" + }, + "outputs": [], + "source": [ + "text = client.call(text, prompt_version=\"deepseek_TN\")\n", + "text" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "qNhCJG4VbSrm" + }, + "outputs": [], + "source": [ + "wav = chat.infer(text)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Wq1XQHmFRQI3" + }, + "outputs": [], + "source": [ + "Audio(wav[0], rate=24_000, autoplay=True)" + ] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "collapsed_sections": [ + "bAUs0rGQbSrh" + ], + "gpuType": "T4", + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.8" + } + }, + "nbformat": 4, + "nbformat_minor": 0 } diff --git a/examples/ipynb/example.ipynb b/examples/ipynb/example.ipynb index 368cdff..7c53ef8 100644 --- a/examples/ipynb/example.ipynb +++ b/examples/ipynb/example.ipynb @@ -19,15 +19,16 @@ " os.environ[\"PYTORCH_ENABLE_MPS_FALLBACK\"] = \"1\"\n", "\n", "if not \"root_dir\" in globals():\n", - " now_dir = os.getcwd() # skip examples/ipynb\n", + " now_dir = os.getcwd() # skip examples/ipynb\n", " root_dir = os.path.join(now_dir, \"../../\")\n", " sys.path.append(root_dir)\n", " print(\"init root dir to\", root_dir)\n", "\n", "import torch\n", + "\n", "torch._dynamo.config.cache_size_limit = 64\n", "torch._dynamo.config.suppress_errors = True\n", - "torch.set_float32_matmul_precision('high')\n", + "torch.set_float32_matmul_precision(\"high\")\n", "\n", "import ChatTTS\n", "from tools.logger import get_logger\n", @@ -59,18 +60,18 @@ "except ValueError as e:\n", " logger.error(e)\n", "except:\n", - " logger.warning('Package nemo_text_processing not found!')\n", + " logger.warning(\"Package nemo_text_processing not found!\")\n", " logger.warning(\n", - " 'Run: conda install -c conda-forge pynini=2.1.5 && pip install nemo_text_processing',\n", + " \"Run: conda install -c conda-forge pynini=2.1.5 && pip install nemo_text_processing\",\n", " )\n", "try:\n", " chat.normalizer.register(\"zh\", normalizer_zh_tn())\n", "except ValueError as e:\n", " logger.error(e)\n", "except:\n", - " logger.warning('Package WeTextProcessing not found!')\n", + " logger.warning(\"Package WeTextProcessing not found!\")\n", " logger.warning(\n", - " 'Run: conda install -c conda-forge pynini=2.1.5 && pip install WeTextProcessing',\n", + " \"Run: conda install -c conda-forge pynini=2.1.5 && pip install WeTextProcessing\",\n", " )" ] }, @@ -95,7 +96,7 @@ "outputs": [], "source": [ "# use force_redownload=True if the weights have been updated.\n", - "chat.load(source='huggingface', force_redownload=True)" + "chat.load(source=\"huggingface\", force_redownload=True)" ] }, { @@ -129,7 +130,7 @@ "outputs": [], "source": [ "# write the model path into custom_path\n", - "chat.load(source='custom', custom_path='YOUR CUSTOM PATH')" + "chat.load(source=\"custom\", custom_path=\"YOUR CUSTOM PATH\")" ] }, { @@ -168,8 +169,11 @@ "metadata": {}, "outputs": [], "source": [ - "texts = [\"So we found being competitive and collaborative was a huge way of staying motivated towards our goals, so one person to call when you fall off, one person who gets you back on then one person to actually do the activity with.\",]*3 \\\n", - " + [\"我觉得像我们这些写程序的人，他，我觉得多多少少可能会对开源有一种情怀在吧我觉得开源是一个很好的形式。现在其实最先进的技术掌握在一些公司的手里的话，就他们并不会轻易的开放给所有的人用。\"]*3 \n", + "texts = [\n", + " \"So we found being competitive and collaborative was a huge way of staying motivated towards our goals, so one person to call when you fall off, one person who gets you back on then one person to actually do the activity with.\",\n", + "] * 3 + [\n", + " \"我觉得像我们这些写程序的人，他，我觉得多多少少可能会对开源有一种情怀在吧我觉得开源是一个很好的形式。现在其实最先进的技术掌握在一些公司的手里的话，就他们并不会轻易的开放给所有的人用。\"\n", + "] * 3\n", "\n", "wavs = chat.infer(texts)" ] @@ -206,15 +210,18 @@ "outputs": [], "source": [ "params_infer_code = ChatTTS.Chat.InferCodeParams(\n", - " prompt='[speed_5]',\n", - " temperature=.3,\n", + " prompt=\"[speed_5]\",\n", + " temperature=0.3,\n", ")\n", "params_refine_text = ChatTTS.Chat.RefineTextParams(\n", - " prompt='[oral_2][laugh_0][break_6]',\n", + " prompt=\"[oral_2][laugh_0][break_6]\",\n", ")\n", "\n", - "wav = chat.infer('四川美食可多了，有麻辣火锅、宫保鸡丁、麻婆豆腐、担担面、回锅肉、夫妻肺片等，每样都让人垂涎三尺。', \\\n", - " params_refine_text=params_refine_text, params_infer_code=params_infer_code)" + "wav = chat.infer(\n", + " \"四川美食可多了，有麻辣火锅、宫保鸡丁、麻婆豆腐、担担面、回锅肉、夫妻肺片等，每样都让人垂涎三尺。\",\n", + " params_refine_text=params_refine_text,\n", + " params_infer_code=params_infer_code,\n", + ")" ] }, { @@ -244,8 +251,11 @@ " spk_emb=rand_spk,\n", ")\n", "\n", - "wav = chat.infer('四川美食确实以辣闻名，但也有不辣的选择。比如甜水面、赖汤圆、蛋烘糕、叶儿粑等，这些小吃口味温和，甜而不腻，也很受欢迎。', \\\n", - " params_refine_text=params_refine_text, params_infer_code=params_infer_code)" + "wav = chat.infer(\n", + " \"四川美食确实以辣闻名，但也有不辣的选择。比如甜水面、赖汤圆、蛋烘糕、叶儿粑等，这些小吃口味温和，甜而不腻，也很受欢迎。\",\n", + " params_refine_text=params_refine_text,\n", + " params_infer_code=params_infer_code,\n", + ")" ] }, { @@ -308,10 +318,10 @@ "source": [ "from tools.llm import ChatOpenAI\n", "\n", - "API_KEY = ''\n", - "client = ChatOpenAI(api_key=API_KEY,\n", - " base_url=\"https://api.deepseek.com\",\n", - " model=\"deepseek-chat\")" + "API_KEY = \"\"\n", + "client = ChatOpenAI(\n", + " api_key=API_KEY, base_url=\"https://api.deepseek.com\", model=\"deepseek-chat\"\n", + ")" ] }, { @@ -320,7 +330,7 @@ "metadata": {}, "outputs": [], "source": [ - "user_question = '四川有哪些好吃的美食呢?'" + "user_question = \"四川有哪些好吃的美食呢?\"" ] }, { @@ -329,7 +339,7 @@ "metadata": {}, "outputs": [], "source": [ - "text = client.call(user_question, prompt_version = 'deepseek')\n", + "text = client.call(user_question, prompt_version=\"deepseek\")\n", "text" ] }, @@ -339,7 +349,7 @@ "metadata": {}, "outputs": [], "source": [ - "text = client.call(text, prompt_version = 'deepseek_TN')\n", + "text = client.call(text, prompt_version=\"deepseek_TN\")\n", "text" ] }, diff --git a/examples/web/funcs.py b/examples/web/funcs.py index a49be4b..5076c61 100644 --- a/examples/web/funcs.py +++ b/examples/web/funcs.py @@ -7,12 +7,14 @@ import numpy as np from tools.audio import unsafe_float_to_int16 from tools.logger import get_logger + logger = get_logger(" WebUI ") from tools.seeder import TorchSeedContext from tools.normalizer import normalizer_en_nemo_text, normalizer_zh_tn import ChatTTS + chat = ChatTTS.Chat(get_logger("ChatTTS")) custom_path: Optional[str] = None @@ -33,19 +35,24 @@ voices = { "Timbre9": {"seed": 9999}, } + def generate_seed(): return gr.update(value=random.randint(1, 100000000)) + # 返回选择音色对应的seed def on_voice_change(vocie_selection): - return voices.get(vocie_selection)['seed'] + return voices.get(vocie_selection)["seed"] + def load_chat(cust_path: Optional[str], coef: Optional[str]) -> bool: if cust_path == None: - ret = chat.load(coef=coef, compile=sys.platform != 'win32') + ret = chat.load(coef=coef, compile=sys.platform != "win32") else: - logger.info('local model path: %s', cust_path) - ret = chat.load('custom', custom_path=cust_path, coef=coef, compile=sys.platform != 'win32') + logger.info("local model path: %s", cust_path) + ret = chat.load( + "custom", custom_path=cust_path, coef=coef, compile=sys.platform != "win32" + ) global custom_path custom_path = cust_path if ret: @@ -54,21 +61,22 @@ def load_chat(cust_path: Optional[str], coef: Optional[str]) -> bool: except ValueError as e: logger.error(e) except: - logger.warning('Package nemo_text_processing not found!') + logger.warning("Package nemo_text_processing not found!") logger.warning( - 'Run: conda install -c conda-forge pynini=2.1.5 && pip install nemo_text_processing', + "Run: conda install -c conda-forge pynini=2.1.5 && pip install nemo_text_processing", ) try: chat.normalizer.register("zh", normalizer_zh_tn()) except ValueError as e: logger.error(e) except: - logger.warning('Package WeTextProcessing not found!') + logger.warning("Package WeTextProcessing not found!") logger.warning( - 'Run: conda install -c conda-forge pynini=2.1.5 && pip install WeTextProcessing', + "Run: conda install -c conda-forge pynini=2.1.5 && pip install WeTextProcessing", ) return ret + def reload_chat(coef: Optional[str]) -> str: chat.unload() gr.Info("Model unloaded.") @@ -85,15 +93,23 @@ def reload_chat(coef: Optional[str]) -> str: gr.Info("Reload succeess.") return chat.coef -def set_generate_buttons(generate_button, interrupt_button, is_reset=False): - return gr.update(value=generate_button, visible=is_reset, interactive=is_reset), gr.update(value=interrupt_button, visible=not is_reset, interactive=not is_reset) -def refine_text(text, text_seed_input, refine_text_flag, generate_button, interrupt_button): +def set_generate_buttons(generate_button, interrupt_button, is_reset=False): + return gr.update( + value=generate_button, visible=is_reset, interactive=is_reset + ), gr.update(value=interrupt_button, visible=not is_reset, interactive=not is_reset) + + +def refine_text( + text, text_seed_input, refine_text_flag, generate_button, interrupt_button +): global chat, has_interrupted has_interrupted = False if not refine_text_flag: - return text, *set_generate_buttons(generate_button, interrupt_button, is_reset=True) + return text, *set_generate_buttons( + generate_button, interrupt_button, is_reset=True + ) with TorchSeedContext(text_seed_input): text = chat.infer( @@ -101,15 +117,20 @@ def refine_text(text, text_seed_input, refine_text_flag, generate_button, interr skip_refine_text=False, refine_text_only=True, ) - return text[0] if isinstance(text, list) else text, *set_generate_buttons(generate_button, interrupt_button, is_reset=True) + return text[0] if isinstance(text, list) else text, *set_generate_buttons( + generate_button, interrupt_button, is_reset=True + ) + def text_output_listener(generate_button, interrupt_button): return set_generate_buttons(generate_button, interrupt_button) + def generate_audio(text, temperature, top_P, top_K, audio_seed_input, stream): global chat, has_interrupted - if not text or text == "𝕃𝕠𝕒𝕕𝕚𝕟𝕘..." or has_interrupted: return None + if not text or text == "𝕃𝕠𝕒𝕕𝕚𝕟𝕘..." or has_interrupted: + return None with TorchSeedContext(audio_seed_input): rand_spk = chat.sample_random_speaker() @@ -138,16 +159,19 @@ def generate_audio(text, temperature, top_P, top_K, audio_seed_input, stream): yield 24000, unsafe_float_to_int16(np.array(wav[0]).flatten()) + def interrupt_generate(): global chat, has_interrupted has_interrupted = True chat.interrupt() + def set_buttons_after_generate(generate_button, interrupt_button, audio_output): global has_interrupted return set_generate_buttons( - generate_button, interrupt_button, + generate_button, + interrupt_button, audio_output is not None or has_interrupted, ) diff --git a/examples/web/webui.py b/examples/web/webui.py index ad24ce5..a200ff6 100644 --- a/examples/web/webui.py +++ b/examples/web/webui.py @@ -12,6 +12,7 @@ import gradio as gr from examples.web.funcs import * + def main(): with gr.Blocks() as demo: @@ -20,53 +21,91 @@ def main(): gr.Markdown("- **HuggingFace Repo**: https://huggingface.co/2Noise/ChatTTS") default_text = "四川美食确实以辣闻名，但也有不辣的选择。比如甜水面、赖汤圆、蛋烘糕、叶儿粑等，这些小吃口味温和，甜而不腻，也很受欢迎。" - text_input = gr.Textbox(label="Input Text", lines=4, placeholder="Please Input Text...", value=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", interactive=True) - 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) + temperature_slider = gr.Slider( + minimum=0.00001, + maximum=1.0, + step=0.00001, + value=0.3, + label="Audio temperature", + interactive=True, + ) + 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='Default') + voice_selection = gr.Dropdown( + label="Timbre", choices=voices.keys(), value="Default" + ) 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") - + with gr.Row(): dvae_coef_text = gr.Textbox( - label="DVAE Coefficient", max_lines=3, show_copy_button=True, scale=4, + label="DVAE Coefficient", + max_lines=3, + show_copy_button=True, + scale=4, ) reload_chat_button = gr.Button("Reload", scale=1) with gr.Row(): auto_play_checkbox = gr.Checkbox(label="Auto Play", value=False, scale=1) - stream_mode_checkbox = gr.Checkbox(label="Stream Mode", value=False, scale=1) + stream_mode_checkbox = gr.Checkbox( + label="Stream Mode", value=False, scale=1 + ) generate_button = gr.Button("Generate", scale=2, variant="primary") - interrupt_button = gr.Button("Interrupt", scale=2, variant="stop", visible=False, interactive=False) + interrupt_button = gr.Button( + "Interrupt", scale=2, variant="stop", visible=False, interactive=False + ) text_output = gr.Textbox(label="Output Text", interactive=False) # 使用Gradio的回调功能来更新数值输入框 - voice_selection.change(fn=on_voice_change, inputs=voice_selection, outputs=audio_seed_input) + voice_selection.change( + fn=on_voice_change, inputs=voice_selection, outputs=audio_seed_input + ) - generate_audio_seed.click(generate_seed, - inputs=[], - outputs=audio_seed_input) + generate_audio_seed.click(generate_seed, inputs=[], outputs=audio_seed_input) + + generate_text_seed.click(generate_seed, inputs=[], outputs=text_seed_input) + + reload_chat_button.click( + reload_chat, inputs=dvae_coef_text, outputs=dvae_coef_text + ) - generate_text_seed.click(generate_seed, - inputs=[], - outputs=text_seed_input) - - reload_chat_button.click(reload_chat, inputs=dvae_coef_text, outputs=dvae_coef_text) - generate_button.click(fn=lambda: "𝕃𝕠𝕒𝕕𝕚𝕟𝕘...", outputs=text_output) - generate_button.click(refine_text, - inputs=[text_input, text_seed_input, refine_text_checkbox, generate_button, interrupt_button], - outputs=[text_output, generate_button, interrupt_button]) - + generate_button.click( + refine_text, + inputs=[ + text_input, + text_seed_input, + refine_text_checkbox, + generate_button, + interrupt_button, + ], + outputs=[text_output, generate_button, interrupt_button], + ) + interrupt_button.click(interrupt_generate) @gr.render(inputs=[auto_play_checkbox, stream_mode_checkbox]) @@ -79,26 +118,81 @@ def main(): interactive=False, show_label=True, ) - text_output.change(text_output_listener, inputs=[generate_button, interrupt_button], outputs=[generate_button, interrupt_button]) - text_output.change(generate_audio, - inputs=[text_output, temperature_slider, top_p_slider, top_k_slider, audio_seed_input, stream_mode_checkbox], - outputs=audio_output).then(fn=set_buttons_after_generate, inputs=[generate_button, interrupt_button, audio_output], outputs=[generate_button, interrupt_button]) + text_output.change( + text_output_listener, + inputs=[generate_button, interrupt_button], + outputs=[generate_button, interrupt_button], + ) + text_output.change( + generate_audio, + inputs=[ + text_output, + temperature_slider, + top_p_slider, + top_k_slider, + audio_seed_input, + stream_mode_checkbox, + ], + outputs=audio_output, + ).then( + fn=set_buttons_after_generate, + inputs=[generate_button, interrupt_button, audio_output], + outputs=[generate_button, interrupt_button], + ) gr.Examples( examples=[ - ["四川美食确实以辣闻名，但也有不辣的选择。比如甜水面、赖汤圆、蛋烘糕、叶儿粑等，这些小吃口味温和，甜而不腻，也很受欢迎。", 0.3, 0.7, 20, 2, 42, True], - ["What is [uv_break]your favorite english food?[laugh][lbreak]", 0.5, 0.5, 10, 245, 531, False], - ["chat T T S is a text to speech model designed for dialogue applications. [uv_break]it supports mixed language input [uv_break]and offers multi speaker capabilities with precise control over prosodic elements [laugh]like like [uv_break]laughter[laugh], [uv_break]pauses, [uv_break]and intonation. [uv_break]it delivers natural and expressive speech,[uv_break]so please[uv_break] use the project responsibly at your own risk.[uv_break]", 0.2, 0.6, 15, 67, 165, False], + [ + "四川美食确实以辣闻名，但也有不辣的选择。比如甜水面、赖汤圆、蛋烘糕、叶儿粑等，这些小吃口味温和，甜而不腻，也很受欢迎。", + 0.3, + 0.7, + 20, + 2, + 42, + True, + ], + [ + "What is [uv_break]your favorite english food?[laugh][lbreak]", + 0.5, + 0.5, + 10, + 245, + 531, + False, + ], + [ + "chat T T S is a text to speech model designed for dialogue applications. [uv_break]it supports mixed language input [uv_break]and offers multi speaker capabilities with precise control over prosodic elements [laugh]like like [uv_break]laughter[laugh], [uv_break]pauses, [uv_break]and intonation. [uv_break]it delivers natural and expressive speech,[uv_break]so please[uv_break] use the project responsibly at your own risk.[uv_break]", + 0.2, + 0.6, + 15, + 67, + 165, + False, + ], + ], + inputs=[ + text_input, + temperature_slider, + top_p_slider, + top_k_slider, + audio_seed_input, + text_seed_input, + refine_text_checkbox, ], - inputs=[text_input, temperature_slider, top_p_slider, top_k_slider, audio_seed_input, text_seed_input, refine_text_checkbox], ) - - 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('--root_path', type=str, default=None, help='root path') - parser.add_argument('--custom_path', type=str, default=None, help='custom model path') - parser.add_argument('--coef', type=str, default=None, help='custom dvae coefficient') + + 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("--root_path", type=str, default=None, help="root path") + parser.add_argument( + "--custom_path", type=str, default=None, help="custom model path" + ) + parser.add_argument( + "--coef", type=str, default=None, help="custom dvae coefficient" + ) args = parser.parse_args() logger.info("loading ChatTTS model...") @@ -111,8 +205,13 @@ def main(): dvae_coef_text.value = chat.coef - demo.launch(server_name=args.server_name, server_port=args.server_port, root_path=args.root_path, inbrowser=True) + demo.launch( + server_name=args.server_name, + server_port=args.server_port, + root_path=args.root_path, + inbrowser=True, + ) -if __name__ == '__main__': +if __name__ == "__main__": main() diff --git a/setup.py b/setup.py index 59ad526..c9153f1 100644 --- a/setup.py +++ b/setup.py @@ -2,32 +2,32 @@ import os from setuptools import setup, find_packages setup( - name='chattts', + name="chattts", version=os.environ.get("CHTTS_VER", "develop"), - description='A generative speech model for daily dialogue', - long_description=open('README.md').read(), - long_description_content_type='text/markdown', - author='2noise', - author_email='open-source@2noise.com', - maintainer='fumiama', - url='https://github.com/2noise/ChatTTS', + description="A generative speech model for daily dialogue", + long_description=open("README.md").read(), + long_description_content_type="text/markdown", + author="2noise", + author_email="open-source@2noise.com", + maintainer="fumiama", + url="https://github.com/2noise/ChatTTS", packages=find_packages(include=["ChatTTS", "ChatTTS.*"]), package_data={ - 'ChatTTS.res': ['homophones_map.json', 'sha256_map.json'], + "ChatTTS.res": ["homophones_map.json", "sha256_map.json"], }, license="CC BY-NC 4.0", install_requires=[ - 'numba', - 'numpy<2.0.0', - 'omegaconf>=2.3.0', - 'pybase16384', - 'torch>=2.1.0', - 'tqdm', - 'transformers>=4.41.1', - 'vector_quantize_pytorch', - 'vocos', + "numba", + "numpy<2.0.0", + "omegaconf>=2.3.0", + "pybase16384", + "torch>=2.1.0", + "tqdm", + "transformers>=4.41.1", + "vector_quantize_pytorch", + "vocos", ], - platforms='any', + platforms="any", classifiers=[ "Programming Language :: Python :: 3", "Operating System :: OS Independent", diff --git a/tools/audio/np.py b/tools/audio/np.py index 9812f78..b40c7f7 100644 --- a/tools/audio/np.py +++ b/tools/audio/np.py @@ -1,6 +1,7 @@ import numpy as np from numba import jit + @jit def unsafe_float_to_int16(audio: np.ndarray) -> np.ndarray: """ diff --git a/tools/llm/llm.py b/tools/llm/llm.py index 352e390..fd6af0f 100644 --- a/tools/llm/llm.py +++ b/tools/llm/llm.py @@ -1,39 +1,74 @@ from openai import OpenAI - + prompt_dict = { - 'kimi': [ {"role": "system", "content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文和英文的对话。"}, - {"role": "user", "content": "你好，请注意你现在生成的文字要按照人日常生活的口吻，你的回复将会后续用TTS模型转为语音，并且请把回答控制在100字以内。并且标点符号仅包含逗号和句号，将数字等转为文字回答。"}, - {"role": "assistant", "content": "好的，我现在生成的文字将按照人日常生活的口吻， 并且我会把回答控制在一百字以内, 标点符号仅包含逗号和句号，将阿拉伯数字等转为中文文字回答。下面请开始对话。"},], - 'deepseek': [ + "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": "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": "你好，现在我们在处理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": "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": "罗森宣布将于七月二十四日退市，在华门店超过六千家。"}, - ], -} - + { + "role": "assistant", + "content": "罗森宣布将于七月二十四日退市，在华门店超过六千家。", + }, + ], +} + + class ChatOpenAI: def __init__(self, api_key, base_url, model): - self.client = OpenAI( - api_key = api_key, - base_url = base_url, + 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): - + + 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, + model=self.model, + messages=prompt_dict[prompt_version] + + [ + {"role": "user", "content": user_question}, + ], + temperature=temperature, **kwargs ) return completion.choices[0].message.content diff --git a/tools/logger/log.py b/tools/logger/log.py index 603a7d5..9321700 100644 --- a/tools/logger/log.py +++ b/tools/logger/log.py @@ -11,11 +11,11 @@ logging.getLogger("NeMo-text-processing").setLevel(logging.WARNING) colorCodePanic = "\x1b[1;31m" colorCodeFatal = "\x1b[1;31m" colorCodeError = "\x1b[31m" -colorCodeWarn = "\x1b[33m" -colorCodeInfo = "\x1b[37m" +colorCodeWarn = "\x1b[33m" +colorCodeInfo = "\x1b[37m" colorCodeDebug = "\x1b[32m" colorCodeTrace = "\x1b[36m" -colorReset = "\x1b[0m" +colorReset = "\x1b[0m" log_level_color_code = { logging.DEBUG: colorCodeDebug, @@ -33,6 +33,7 @@ log_level_msg_str = { logging.FATAL: "FATL", } + class Formatter(logging.Formatter): def __init__(self, color=platform.system().lower() != "windows"): # https://stackoverflow.com/questions/2720319/python-figure-out-local-timezone @@ -40,7 +41,7 @@ class Formatter(logging.Formatter): self.color = color def format(self, record: logging.LogRecord): - logstr = "[" + datetime.now(self.tz).strftime('%z %Y%m%d %H:%M:%S') + "] [" + logstr = "[" + datetime.now(self.tz).strftime("%z %Y%m%d %H:%M:%S") + "] [" if self.color: logstr += log_level_color_code.get(record.levelno, colorCodeInfo) logstr += log_level_msg_str.get(record.levelno, record.levelname) @@ -51,7 +52,7 @@ class Formatter(logging.Formatter): return logstr -def get_logger(name: str, lv = logging.INFO, remove_exist=False, format_root=False): +def get_logger(name: str, lv=logging.INFO, remove_exist=False, format_root=False): logger = logging.getLogger(name) logger.setLevel(lv) if remove_exist and logger.hasHandlers(): diff --git a/tools/normalizer/en.py b/tools/normalizer/en.py index ce092c8..980043b 100644 --- a/tools/normalizer/en.py +++ b/tools/normalizer/en.py @@ -1,9 +1,12 @@ from typing import Callable from functools import partial + def normalizer_en_nemo_text() -> Callable[[str], str]: from nemo_text_processing.text_normalization.normalize import Normalizer + return partial( - Normalizer(input_case='cased', lang="en").normalize, - verbose=False, punct_post_process=True, + Normalizer(input_case="cased", lang="en").normalize, + verbose=False, + punct_post_process=True, ) diff --git a/tools/normalizer/zh.py b/tools/normalizer/zh.py index d3087a4..b513d19 100644 --- a/tools/normalizer/zh.py +++ b/tools/normalizer/zh.py @@ -1,5 +1,7 @@ from typing import Callable + def normalizer_zh_tn() -> Callable[[str], str]: from tn.chinese.normalizer import Normalizer + return Normalizer().normalize diff --git a/tools/seeder/ctx.py b/tools/seeder/ctx.py index ab510f4..30bf4c9 100644 --- a/tools/seeder/ctx.py +++ b/tools/seeder/ctx.py @@ -1,5 +1,6 @@ import torch + class TorchSeedContext: def __init__(self, seed): self.seed = seed