From 4dd1f88073838d8713d8cba28c4d6e922b74ac14 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Sun, 23 Jun 2024 17:17:33 +0900
Subject: [PATCH 1/9] optimize: core & model

### all
- reduce redundant code
- add `_` to some inner calls
### core
- rename `check_model` to `has_loaded`
- new api `download_models`
- move vocos out of models dict (further target: remove the dict `pretrain_models`)
- new api `unload`
- new api `decode_to_wavs`
### model
- fix grad mul issue in dvae
- optimize progress bar display
- optimize tensor operations
---
 ChatTTS/core.py              | 193 +++++++++++++++++++----------------
 ChatTTS/model/dvae.py        |  27 ++---
 ChatTTS/model/gpt.py         | 114 ++++++++++++---------
 ChatTTS/utils/io.py          |   7 +-
 examples/ipynb/colab.ipynb   |  16 +++
 examples/ipynb/example.ipynb |  16 +++
 tools/logger/log.py          |   4 +-
 7 files changed, 223 insertions(+), 154 deletions(-)

diff --git a/ChatTTS/core.py b/ChatTTS/core.py
index dc3f658..2358109 100644
--- a/ChatTTS/core.py
+++ b/ChatTTS/core.py
@@ -3,7 +3,7 @@ import json
 import logging
 import tempfile
 from functools import partial
-from typing import Literal, Optional
+from typing import Literal, Optional, List
 
 import torch
 from omegaconf import OmegaConf
@@ -28,25 +28,66 @@ class Chat:
         self.logger = logger
         set_utils_logger(logger)
         
-    def check_model(self, level = logging.INFO, use_decoder = False):
+    def has_loaded(self, use_decoder = False):
         not_finish = False
-        check_list = ['vocos', 'gpt', 'tokenizer']
+        check_list = ['gpt', 'tokenizer']
         
         if use_decoder:
             check_list.append('decoder')
         else:
             check_list.append('dvae')
-            
+
         for module in check_list:
             if module not in self.pretrain_models:
-                self.logger.log(logging.WARNING, f'{module} not initialized.')
+                self.logger.warn(f'{module} not initialized.')
                 not_finish = True
-                
+
+        if not hasattr(self, "_vocos_decode") or not hasattr(self, "vocos"):
+            self.logger.warn('vocos not initialized.')
+            not_finish = True
+
         if not not_finish:
-            self.logger.log(level, f'All initialized.')
+            self.logger.info('all models has been initialized.')
 
         return not not_finish
 
+    def download_models(
+        self,
+        source: Literal['huggingface', 'local', 'custom']='local',
+        force_redownload=False,
+        custom_path: Optional[torch.serialization.FILE_LIKE]=None,
+    ) -> Optional[str]:
+        if source == 'local':
+            download_path = os.getcwd()
+            if not check_all_assets(update=True) or force_redownload:
+                with tempfile.TemporaryDirectory() as tmp:
+                    download_all_assets(tmpdir=tmp)
+                if not check_all_assets(update=False):
+                    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"))
+            try:
+                download_path = get_latest_modified_file(os.path.join(hf_home, 'hub/models--2Noise--ChatTTS/snapshots'))
+            except:
+                download_path = None
+            if download_path is None or force_redownload: 
+                self.logger.log(logging.INFO, f'Download from HF: https://huggingface.co/2Noise/ChatTTS')
+                try:
+                    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}')
+            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}')
+            download_path = custom_path
+        
+        return download_path
+
     def load_models(
         self,
         source: Literal['huggingface', 'local', 'custom']='local',
@@ -55,36 +96,15 @@ class Chat:
         custom_path: Optional[torch.serialization.FILE_LIKE]=None,
         device: Optional[torch.device] = None,
         coef: Optional[torch.Tensor] = None,
-    ):
-        if source == 'local':
-            torch.load
-            download_path = os.getcwd()
-            if not check_all_assets(update=True) or force_redownload:
-                with tempfile.TemporaryDirectory() as tmp:
-                    download_all_assets(tmpdir=tmp)
-                if not check_all_assets(update=False):
-                    self.logger.error("counld not satisfy all assets needed.")
-                    return False
-        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'))
-            except:
-                download_path = None
-            if download_path is None or force_redownload: 
-                self.logger.log(logging.INFO, f'Download from HF: https://huggingface.co/2Noise/ChatTTS')
-                download_path = snapshot_download(repo_id="2Noise/ChatTTS", allow_patterns=["*.pt", "*.yaml"])
-            else:
-                self.logger.log(logging.INFO, f'Load from cache: {download_path}')
-        elif source == 'custom':
-            self.logger.log(logging.INFO, f'Load from local: {custom_path}')
-            download_path = custom_path
-
+    ) -> bool:
+        download_path = self.download_models(source, force_redownload, custom_path)
+        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()},
         )
-        
+
     def _load(
         self, 
         vocos_config_path: str = None, 
@@ -112,9 +132,17 @@ class Chat:
             ).eval()
             assert vocos_ckpt_path, 'vocos_ckpt_path should not be None'
             vocos.load_state_dict(torch.load(vocos_ckpt_path))
-            self.pretrain_models['vocos'] = vocos
+            self.vocos = vocos
+            if "mps" in str(self.device):
+                self._vocos_decode = lambda spec: self.vocos.decode(
+                    spec.cpu()
+                ).cpu().numpy()
+            else:
+                self._vocos_decode = 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()
@@ -157,8 +185,13 @@ class Chat:
         
         self.coef = coef
 
-        return self.check_model()
+        return self.has_loaded()
     
+    def unload(self):
+        logger = self.logger
+        del_all(self)
+        self.__init__(logger)
+
     def _infer(
         self, 
         text, 
@@ -173,14 +206,14 @@ class Chat:
         do_homophone_replacement=True
     ):
         
-        assert self.check_model(use_decoder=use_decoder)
+        assert self.has_loaded(use_decoder=use_decoder)
         
         if not isinstance(text, list): 
             text = [text]
         if do_text_normalization:
             for i, t in enumerate(text):
                 _lang = detect_language(t) if lang is None else lang
-                if self.init_normalizer(_lang):
+                if self._init_normalizer(_lang):
                     text[i] = self.normalizer[_lang](t)
                     if _lang == 'zh':
                         text[i] = apply_half2full_map(text[i])
@@ -189,7 +222,7 @@ class Chat:
             if len(invalid_characters):
                 self.logger.log(logging.WARNING, f'Invalid characters found! : {invalid_characters}')
                 text[i] = apply_character_map(t)
-            if do_homophone_replacement and self.init_homophones_replacer():
+            if do_homophone_replacement and self._init_homophones_replacer():
                 text[i], replaced_words = self.homophones_replacer.replace(text[i])
                 if replaced_words:
                     repl_res = ', '.join([f'{_[0]}->{_[1]}' for _ in replaced_words])
@@ -205,64 +238,25 @@ class Chat:
             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)
-            del_all(refined)
+            refined.destroy()
             if refine_text_only:
                 yield text
                 return
 
         text = [params_infer_code.get('prompt', '') + i for i in text]
         params_infer_code.pop('prompt', '')
-        result_gen = infer_code(
+
+        length = [0 for _ in range(len(text))]
+        for result in infer_code(
             self.pretrain_models,
             text,
             device=self.device,
             **params_infer_code,
             return_hidden=use_decoder,
             stream=stream,
-        )
-        if use_decoder:
-            docoder_name = 'decoder'
-        else:
-            docoder_name = 'dvae'
-        if "mps" in str(self.device):
-            vocos_decode = lambda spec: [self.pretrain_models['vocos'].decode(
-                i.cpu()
-            ).cpu().numpy() for i in spec]
-        else:
-            vocos_decode = lambda spec: [self.pretrain_models['vocos'].decode(
-                i
-            ).cpu().numpy() for i in spec]
-        if stream:
-
-            length = 0
-            for result in result_gen:
-                x = result.hiddens if use_decoder else result.ids
-                assert len(x) == 1
-                chunk_data = x[0]
-                start_seek = length
-                length = len(chunk_data)
-                self.logger.debug(f'{start_seek=} total len: {length}, new len: {length - start_seek = }')
-                chunk_data = chunk_data[start_seek:]
-                if not len(chunk_data):
-                    continue
-                self.logger.debug(f'new hidden {len(chunk_data)=}')
-                mel_spec = [self.pretrain_models[docoder_name](i[None].permute(0,2,1).to(self.device)) for i in [chunk_data]]
-                del_all(result)
-                del chunk_data
-                del_all(x)
-                wav = vocos_decode(mel_spec)
-                del_all(mel_spec)
-                self.logger.debug(f'yield wav chunk {len(wav[0])=} {len(wav[0][0])=}')
-                yield wav
-            return
-        result = next(result_gen)
-        x = result.hiddens if use_decoder else result.ids
-        mel_spec = [self.pretrain_models[docoder_name](i[None].permute(0,2,1).to(self.device)) for i in x]
-        del_all(result)
-        del_all(x)
-        wav = vocos_decode(mel_spec)
-        del_all(mel_spec)
-        yield wav
+        ):
+            wav = self.decode_to_wavs(result, length, use_decoder)
+            yield wav
 
     def infer(
         self, 
@@ -299,8 +293,31 @@ class Chat:
         dim = self.pretrain_models['gpt'].gpt.layers[0].mlp.gate_proj.in_features
         std, mean = self.pretrain_models['spk_stat'].chunk(2)
         return torch.randn(dim, device=std.device) * std + mean
-    
-    def init_normalizer(self, lang) -> 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 = []
+        for i, chunk_data in enumerate(x):
+            start_seek = start_seeks[i]
+            length = len(chunk_data)
+            if length <= start_seek:
+                wavs.append(None)
+                continue
+            start_seeks[i] = length
+            chunk_data = chunk_data[start_seek:]
+            if use_decoder:
+                decoder = self.pretrain_models['decoder']
+            else:
+                decoder = self.pretrain_models['dvae']
+            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 _init_normalizer(self, lang) -> bool:
 
         if lang in self.normalizer:
             return True
@@ -335,7 +352,7 @@ class Chat:
                 )
         return False
 
-    def init_homophones_replacer(self):
+    def _init_homophones_replacer(self):
         if self.homophones_replacer:
             return True
         else:
diff --git a/ChatTTS/model/dvae.py b/ChatTTS/model/dvae.py
index 8e22304..4090d73 100644
--- a/ChatTTS/model/dvae.py
+++ b/ChatTTS/model/dvae.py
@@ -164,20 +164,21 @@ class DVAE(nn.Module):
         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():
 
-        if self.vq_layer is not None:
-            vq_feats = self.vq_layer._embed(inp)
-        else:
-            vq_feats = inp.detach().clone()
+            if self.vq_layer is not None:
+                vq_feats = self.vq_layer._embed(inp)
+            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(
-                input=vq_feats.transpose_(1, 2),
-            ).transpose_(1, 2),
-        )
+            dec_out = self.out_conv(
+                self.decoder(
+                    input=vq_feats.transpose_(1, 2),
+                ).transpose_(1, 2),
+            )
 
-        return torch.mul(dec_out, self.coef, out=dec_out)
+            return torch.mul(dec_out, self.coef, out=dec_out)
diff --git a/ChatTTS/model/gpt.py b/ChatTTS/model/gpt.py
index 3a26a3d..32bb952 100644
--- a/ChatTTS/model/gpt.py
+++ b/ChatTTS/model/gpt.py
@@ -1,20 +1,19 @@
 import os
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
+from dataclasses import dataclass
 import logging
 from typing import Union, List, Optional, Tuple
-from dataclasses import dataclass
 
-from tqdm import tqdm
-from transformers.cache_utils import Cache
 import omegaconf
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.nn.utils.parametrize as P
 from torch.nn.utils.parametrizations import weight_norm
+from tqdm import tqdm
 from transformers import LlamaModel, LlamaConfig, LogitsWarper
+from transformers.cache_utils import Cache
 from transformers.modeling_outputs import BaseModelOutputWithPast
 
 from ..utils.infer import CustomRepetitionPenaltyLogitsProcessorRepeat
@@ -98,23 +97,22 @@ class GPT(nn.Module):
 
         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
+        text_mask_inv = ~(text_mask.to(self.device_gpt))
         masked_input_ids: torch.Tensor = input_ids[text_mask_inv].to(self.device_gpt)
-        del text_mask_inv
 
         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[text_mask] = emb_text
-        emb[~text_mask] = emb_code.to(emb.dtype)
+        emb[text_mask_inv] = emb_code.to(emb.dtype)
 
-        del emb_text, emb_code
+        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
@@ -129,7 +127,7 @@ class GPT(nn.Module):
             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_inputs_for_generation(
+    def _prepare_generation_inputs(
         self,
         input_ids: torch.Tensor,
         past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]]=None,
@@ -138,7 +136,7 @@ class GPT(nn.Module):
         cache_position: Optional[torch.Tensor]=None,
         position_ids: Optional[torch.Tensor]=None,
         use_cache = True,
-    ) -> GenerationInputs:
+    ) -> _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
@@ -197,7 +195,7 @@ class GPT(nn.Module):
         if has_static_cache:
             past_key_values = None
 
-        model_inputs = self.GenerationInputs(
+        model_inputs = self._GenerationInputs(
             position_ids=position_ids,
             cache_position=cache_position,
             use_cache=use_cache,
@@ -223,7 +221,36 @@ class GPT(nn.Module):
         attentions: List[Optional[Tuple[torch.FloatTensor, ...]]]
         hiddens: List[torch.Tensor]
 
-    
+        def destroy(self):
+            del_all(self.ids)
+            del_all(self.attentions)
+            del_all(self.hiddens)
+
+
+    def _prepare_generation_outputs(
+        self,
+        inputs_ids: torch.Tensor,
+        start_idx: int,
+        end_idx: torch.Tensor,
+        attentions: List[Optional[Tuple[torch.FloatTensor, ...]]],
+        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)]
+        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())]
+
+        return self.GenerationOutputs(
+            ids=inputs_ids,
+            attentions=attentions,
+            hiddens=hiddens,
+        )
+
+
     def generate(
         self, 
         emb: torch.Tensor, 
@@ -256,12 +283,16 @@ class GPT(nn.Module):
             if attention_mask is not None:
                 attention_mask_cache[:, :attention_mask.shape[1]] = attention_mask
 
-            with tqdm(total=max_new_token) as pbar:
+            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}]',
+            ) as pbar:
 
                 past_key_values = None
 
                 for i in range(max_new_token):
-                    model_input = self._prepare_inputs_for_generation(
+                    model_input = self._prepare_generation_inputs(
                         inputs_ids, 
                         past_key_values, 
                         attention_mask_cache[:, :inputs_ids.shape[1]],
@@ -296,7 +327,7 @@ class GPT(nn.Module):
                     past_key_values = outputs.past_key_values
                     del_all(outputs)
                     if return_hidden:
-                        hiddens.append(hidden_states[:, -1])
+                        hiddens.append(hidden_states.narrow(1, -1, 1).squeeze_(1))
 
                     with P.cached():
                         if infer_text:
@@ -353,52 +384,39 @@ class GPT(nn.Module):
                         finish_or = (idx_next == eos_token).any(1)
                         finish |= finish_or
                         del finish_or
-                        inputs_ids = 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 |= finish_or
                         del finish_or
-                        inputs_ids = 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 idx_next
+                    del inputs_ids
+                    inputs_ids = inputs_ids_tmp
+                    del inputs_ids_tmp, idx_next
 
-                    end_idx += (~finish).int().to(end_idx.device)
                     if stream:
-                        if end_idx % 24 and not finish.all():
-                            continue
-                        y_inputs_ids = [inputs_ids[idx, start_idx: start_idx+i] for idx, i in enumerate(end_idx.int())]
-                        y_inputs_ids = [i[:, 0] for i in y_inputs_ids] if infer_text else y_inputs_ids
-                        y_hiddens = []
-                        if return_hidden:
-                            y_hiddens = torch.stack(hiddens, 1)
-                            y_hiddens = [y_hiddens[idx, :i] for idx, i in enumerate(end_idx.int())]
+                        minus_prev_end_index = -end_idx
+                    end_idx += (~finish.to(end_idx.device)).int()
+                    if stream:
+                        if end_idx.all() and (end_idx%24 == 0).any() and torch.add(end_idx, minus_prev_end_index, out=minus_prev_end_index).any():
+                            self.logger.debug("yield stream result, end: %d", end_idx)
+                            yield self._prepare_generation_outputs(
+                                inputs_ids, start_idx, end_idx, attentions, hiddens,
+                                infer_text,
+                            )
+                        del minus_prev_end_index
 
-                        yield self.GenerationOutputs(
-                            ids=y_inputs_ids,
-                            attentions=attentions,
-                            hiddens=y_hiddens,
-                        )
-
-                    if finish.all():
-                        pbar.update(max_new_token-i-1)
-                        break
+                    if finish.all(): break
 
                     pbar.update(1)
 
-            inputs_ids = [inputs_ids[idx, start_idx: start_idx+i] for idx, i in enumerate(end_idx.int())]
-            inputs_ids = [i[:, 0] for i in inputs_ids] if infer_text else inputs_ids
-            
-            if return_hidden:
-                hiddens = torch.stack(hiddens, 1)
-                hiddens = [hiddens[idx, :i] for idx, i in enumerate(end_idx.int())]
-                    
             if not finish.all():
                 self.logger.warn(f'Incomplete result. hit max_new_token: {max_new_token}')
 
             del finish
 
-            yield self.GenerationOutputs(
-                ids=inputs_ids,
-                attentions=attentions,
-                hiddens=hiddens,
+            yield self._prepare_generation_outputs(
+                inputs_ids, start_idx, end_idx, attentions, hiddens,
+                infer_text,
             )
diff --git a/ChatTTS/utils/io.py b/ChatTTS/utils/io.py
index d66c3fa..eeb9871 100644
--- a/ChatTTS/utils/io.py
+++ b/ChatTTS/utils/io.py
@@ -10,7 +10,7 @@ def get_latest_modified_file(directory):
 
     files = [os.path.join(directory, f) for f in os.listdir(directory)] 
     if not files:
-        logger.log(logging.WARNING, f'No files found in the directory: {directory}')
+        logger.log(logging.WARNING, f'no files found in the directory: {directory}')
         return None
     latest_file = max(files, key=os.path.getmtime)
 
@@ -31,11 +31,12 @@ def del_all(d: Union[dict, list]):
             if isinstance(x, dict) or isinstance(x, list) or is_dataclass(x):
                 del_all(x)
             del x
-        return
     elif isinstance(d, list):
         while len(d):
             x = d.pop()
             if isinstance(x, dict) or isinstance(x, list) or is_dataclass(x):
                 del_all(x)
             del x
-        return
+    else:
+        del d
+
diff --git a/examples/ipynb/colab.ipynb b/examples/ipynb/colab.ipynb
index ffd4aa4..589eaac 100644
--- a/examples/ipynb/colab.ipynb
+++ b/examples/ipynb/colab.ipynb
@@ -132,6 +132,22 @@
         "chat.load_models(source='custom', custom_path='YOUR CUSTOM PATH')"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "### You can also unload models to save the memory"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "chat.unload()"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
diff --git a/examples/ipynb/example.ipynb b/examples/ipynb/example.ipynb
index ced24e9..0ef0e6a 100644
--- a/examples/ipynb/example.ipynb
+++ b/examples/ipynb/example.ipynb
@@ -113,6 +113,22 @@
     "chat.load_models(source='custom', custom_path='YOUR CUSTOM PATH')"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### You can also unload models to save the memory"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "chat.unload()"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
diff --git a/tools/logger/log.py b/tools/logger/log.py
index ce08989..7d26ce7 100644
--- a/tools/logger/log.py
+++ b/tools/logger/log.py
@@ -4,6 +4,8 @@ from datetime import datetime, timezone
 
 logging.getLogger("numba").setLevel(logging.WARNING)
 logging.getLogger("httpx").setLevel(logging.WARNING)
+logging.getLogger("wetext-zh_normalizer").setLevel(logging.WARNING)
+logging.getLogger("NeMo-text-processing").setLevel(logging.WARNING)
 
 # from https://github.com/FloatTech/ZeroBot-Plugin/blob/c70766a989698452e60e5e48fb2f802a2444330d/console/console_windows.go#L89-L96
 colorCodePanic = "\x1b[1;31m"
@@ -48,8 +50,6 @@ class Formatter(logging.Formatter):
         logstr += f"] {str(record.name)} | {fn} | {str(record.msg)%record.args}"
         return logstr
 
-for h in logging.root.handlers:
-    h.setFormatter(Formatter())
 
 def get_logger(name: str, lv = logging.INFO):
     logger = logging.getLogger(name)

From 78ee6965b7cd0718d4e94a409d9d9464919a4eb6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Sun, 23 Jun 2024 17:21:23 +0900
Subject: [PATCH 2/9] chore(webui): unify display language

---
 examples/web/funcs.py | 21 ++++++++++-----------
 examples/web/webui.py |  2 +-
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/examples/web/funcs.py b/examples/web/funcs.py
index e0dedb9..ab673ea 100644
--- a/examples/web/funcs.py
+++ b/examples/web/funcs.py
@@ -13,17 +13,16 @@ chat = ChatTTS.Chat(get_logger("ChatTTS"))
 
 # 音色选项：用于预置合适的音色
 voices = {
-    "默认": {"seed": 2},
-    "音色1": {"seed": 1111},
-    "音色2": {"seed": 2222},
-    "音色3": {"seed": 3333},
-    "音色4": {"seed": 4444},
-    "音色5": {"seed": 5555},
-    "音色6": {"seed": 6666},
-    "音色7": {"seed": 7777},
-    "音色8": {"seed": 8888},
-    "音色9": {"seed": 9999},
-    "音色10": {"seed": 11111},
+    "Default": {"seed": 2},
+    "Timbre1": {"seed": 1111},
+    "Timbre2": {"seed": 2222},
+    "Timbre3": {"seed": 3333},
+    "Timbre4": {"seed": 4444},
+    "Timbre5": {"seed": 5555},
+    "Timbre6": {"seed": 6666},
+    "Timbre7": {"seed": 7777},
+    "Timbre8": {"seed": 8888},
+    "Timbre9": {"seed": 9999},
 }
 
 def generate_seed():
diff --git a/examples/web/webui.py b/examples/web/webui.py
index 135a6d6..0d01006 100644
--- a/examples/web/webui.py
+++ b/examples/web/webui.py
@@ -32,7 +32,7 @@ def main():
             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="音色", choices=voices.keys(), value='默认')
+            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")

From fef7084d1cdb4ba6c2b4abfcf0f4b7b05389347a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Sun, 23 Jun 2024 18:04:56 +0900
Subject: [PATCH 3/9] feat(webui): add reload button & dvae coef

---
 ChatTTS/core.py       |  3 +--
 examples/web/funcs.py | 31 ++++++++++++++++++++++++++-----
 examples/web/webui.py | 22 ++++++++++++++++------
 3 files changed, 43 insertions(+), 13 deletions(-)

diff --git a/ChatTTS/core.py b/ChatTTS/core.py
index 2358109..4639c88 100644
--- a/ChatTTS/core.py
+++ b/ChatTTS/core.py
@@ -288,8 +288,7 @@ class Chat:
         else:
             return next(res_gen)
     
-    def sample_random_speaker(self, ):
-        
+    def sample_random_speaker(self):
         dim = self.pretrain_models['gpt'].gpt.layers[0].mlp.gate_proj.in_features
         std, mean = self.pretrain_models['spk_stat'].chunk(2)
         return torch.randn(dim, device=std.device) * std + mean
diff --git a/examples/web/funcs.py b/examples/web/funcs.py
index ab673ea..62044ac 100644
--- a/examples/web/funcs.py
+++ b/examples/web/funcs.py
@@ -1,4 +1,5 @@
 import random
+from typing import Optional
 
 import torch
 import gradio as gr
@@ -11,6 +12,8 @@ logger = get_logger(" WebUI ")
 import ChatTTS
 chat = ChatTTS.Chat(get_logger("ChatTTS"))
 
+custom_path: Optional[str] = None
+
 # 音色选项：用于预置合适的音色
 voices = {
     "Default": {"seed": 2},
@@ -32,13 +35,32 @@ def generate_seed():
 def on_voice_change(vocie_selection):
     return voices.get(vocie_selection)['seed']
 
-def refine_text(text, audio_seed_input, text_seed_input, refine_text_flag):
+def reload_chat(coef: Optional[str]) -> str:
+    global custom_path
+    chat.unload()
+    gr.Info("Model unloaded.")
+    try:
+        if len(coef) != 230:
+            gr.Warning("Ingore invalid DVAE coefficient.")
+            coef = None
+        if custom_path == None:
+            ret = chat.load_models(coef=coef)
+        else:
+            logger.info('local model path: %s', custom_path)
+            ret = chat.load_models('custom', custom_path=custom_path, coef=coef)
+        if not ret:
+            raise gr.Error("Unable to load model.")
+        gr.Info("Reload succeess.")
+        return chat.coef
+    except Exception as e:
+        raise gr.Error(str(e))
+
+def refine_text(text, text_seed_input, refine_text_flag):
     if not refine_text_flag:
         return text
 
     global chat
 
-    torch.manual_seed(audio_seed_input)
     params_refine_text = {'prompt': '[oral_2][laugh_0][break_6]'}
 
     torch.manual_seed(text_seed_input)
@@ -50,7 +72,7 @@ def refine_text(text, audio_seed_input, text_seed_input, refine_text_flag):
                         )
     return text[0] if isinstance(text, list) else text
 
-def generate_audio(text, temperature, top_P, top_K, audio_seed_input, text_seed_input, stream):
+def generate_audio(text, temperature, top_P, top_K, audio_seed_input, stream):
     if not text: return None
 
     global chat
@@ -62,8 +84,7 @@ def generate_audio(text, temperature, top_P, top_K, audio_seed_input, text_seed_
         'temperature': temperature,
         'top_P': top_P,
         'top_K': top_K,
-        }
-    torch.manual_seed(text_seed_input)
+    }
 
     wav = chat.infer(
         text,
diff --git a/examples/web/webui.py b/examples/web/webui.py
index 0d01006..23c76ee 100644
--- a/examples/web/webui.py
+++ b/examples/web/webui.py
@@ -37,6 +37,12 @@ def main():
             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,
+            )
+            reload_chat_button = gr.Button("Reload", scale=1)
 
         with gr.Row():
             auto_play_checkbox = gr.Checkbox(label="Auto Play", value=False, scale=1)
@@ -56,9 +62,11 @@ def main():
                                  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, audio_seed_input, text_seed_input, refine_text_checkbox],
+                              inputs=[text_input, text_seed_input, refine_text_checkbox],
                               outputs=text_output)
 
         @gr.render(inputs=[auto_play_checkbox, stream_mode_checkbox])
@@ -72,14 +80,14 @@ def main():
                 show_label=True,
             )
             text_output.change(generate_audio,
-                                inputs=[text_output, temperature_slider, top_p_slider, top_k_slider, audio_seed_input, text_seed_input, stream_mode_checkbox],
+                                inputs=[text_output, temperature_slider, top_p_slider, top_k_slider, audio_seed_input, stream_mode_checkbox],
                                 outputs=audio_output)
 
         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, True],
-                ["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, 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],
         )
@@ -93,7 +101,7 @@ def main():
 
     logger.info("loading ChatTTS model...")
 
-    global chat
+    global chat, custom_path
 
     if args.custom_path == None:
         ret = chat.load_models()
@@ -106,7 +114,9 @@ def main():
     else:
         logger.error("Models load failed.")
         sys.exit(1)
-
+    
+    custom_path = args.custom_path
+    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)
 

From 93ef96b173d610dd78d56ba0f58c7a95a9e15e21 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Sun, 23 Jun 2024 21:22:23 +0900
Subject: [PATCH 4/9] feat: apply TorchSeedContext

---
 examples/web/funcs.py    | 44 ++++++++++++++++++++--------------------
 tools/seeder/__init__.py |  1 +
 tools/seeder/ctx.py      | 13 ++++++++++++
 3 files changed, 36 insertions(+), 22 deletions(-)
 create mode 100644 tools/seeder/__init__.py
 create mode 100644 tools/seeder/ctx.py

diff --git a/examples/web/funcs.py b/examples/web/funcs.py
index 62044ac..5351a95 100644
--- a/examples/web/funcs.py
+++ b/examples/web/funcs.py
@@ -1,7 +1,6 @@
 import random
 from typing import Optional
 
-import torch
 import gradio as gr
 import numpy as np
 
@@ -9,6 +8,8 @@ from tools.audio import unsafe_float_to_int16
 from tools.logger import get_logger
 logger = get_logger(" WebUI ")
 
+from tools.seeder import TorchSeedContext
+
 import ChatTTS
 chat = ChatTTS.Chat(get_logger("ChatTTS"))
 
@@ -63,13 +64,12 @@ def refine_text(text, text_seed_input, refine_text_flag):
 
     params_refine_text = {'prompt': '[oral_2][laugh_0][break_6]'}
 
-    torch.manual_seed(text_seed_input)
-
-    text = chat.infer(text,
-                        skip_refine_text=False,
-                        refine_text_only=True,
-                        params_refine_text=params_refine_text,
-                        )
+    with TorchSeedContext(text_seed_input):
+        text = chat.infer(text,
+                            skip_refine_text=False,
+                            refine_text_only=True,
+                            params_refine_text=params_refine_text,
+                            )
     return text[0] if isinstance(text, list) else text
 
 def generate_audio(text, temperature, top_P, top_K, audio_seed_input, stream):
@@ -77,21 +77,21 @@ def generate_audio(text, temperature, top_P, top_K, audio_seed_input, stream):
 
     global chat
 
-    torch.manual_seed(audio_seed_input)
-    rand_spk = chat.sample_random_speaker()
-    params_infer_code = {
-        'spk_emb': rand_spk,
-        'temperature': temperature,
-        'top_P': top_P,
-        'top_K': top_K,
-    }
+    with TorchSeedContext(audio_seed_input):
+        rand_spk = chat.sample_random_speaker()
+        params_infer_code = {
+            'spk_emb': rand_spk,
+            'temperature': temperature,
+            'top_P': top_P,
+            'top_K': top_K,
+        }
 
-    wav = chat.infer(
-        text,
-        skip_refine_text=True,
-        params_infer_code=params_infer_code,
-        stream=stream,
-    )
+        wav = chat.infer(
+            text,
+            skip_refine_text=True,
+            params_infer_code=params_infer_code,
+            stream=stream,
+        )
 
     if stream:
         for gen in wav:
diff --git a/tools/seeder/__init__.py b/tools/seeder/__init__.py
new file mode 100644
index 0000000..a7d73d0
--- /dev/null
+++ b/tools/seeder/__init__.py
@@ -0,0 +1 @@
+from .ctx import TorchSeedContext
diff --git a/tools/seeder/ctx.py b/tools/seeder/ctx.py
new file mode 100644
index 0000000..ab510f4
--- /dev/null
+++ b/tools/seeder/ctx.py
@@ -0,0 +1,13 @@
+import torch
+
+class TorchSeedContext:
+    def __init__(self, seed):
+        self.seed = seed
+        self.state = None
+
+    def __enter__(self):
+        self.state = torch.random.get_rng_state()
+        torch.manual_seed(self.seed)
+
+    def __exit__(self, type, value, traceback):
+        torch.random.set_rng_state(self.state)

From 79684e242616995e33ceaa05c311a77bdd18fd6d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Sun, 23 Jun 2024 21:26:23 +0900
Subject: [PATCH 5/9] feat(webui): make `Generate` button be primary

---
 examples/web/webui.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/web/webui.py b/examples/web/webui.py
index 23c76ee..7f185ef 100644
--- a/examples/web/webui.py
+++ b/examples/web/webui.py
@@ -47,7 +47,7 @@ def main():
         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)
-            generate_button = gr.Button("Generate", scale=2)
+            generate_button = gr.Button("Generate", scale=2, variant="primary")
 
         text_output = gr.Textbox(label="Output Text", interactive=False)
 

From 5d2002831645155c1d244bf00407f49a02ae14b0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Sun, 23 Jun 2024 22:39:11 +0900
Subject: [PATCH 6/9] optimize(utils): jit HomophonesReplacer

---
 ChatTTS/core.py        |   8 ++--
 ChatTTS/utils/infer.py | 101 +++++++++++++++++++++++++----------------
 tools/audio/np.py      |   3 +-
 3 files changed, 68 insertions(+), 44 deletions(-)

diff --git a/ChatTTS/core.py b/ChatTTS/core.py
index 4639c88..e3891cd 100644
--- a/ChatTTS/core.py
+++ b/ChatTTS/core.py
@@ -220,7 +220,7 @@ class Chat:
         for i, t in enumerate(text):
             invalid_characters = count_invalid_characters(t)
             if len(invalid_characters):
-                self.logger.log(logging.WARNING, f'Invalid characters found! : {invalid_characters}')
+                self.logger.warn(f'Invalid characters found! : {invalid_characters}')
                 text[i] = apply_character_map(t)
             if do_homophone_replacement and self._init_homophones_replacer():
                 text[i], replaced_words = self.homophones_replacer.replace(text[i])
@@ -357,10 +357,10 @@ class Chat:
         else:
             try:
                 self.homophones_replacer = HomophonesReplacer(os.path.join(os.path.dirname(__file__), 'res', 'homophones_map.json'))
-                self.logger.log(logging.INFO, 'homophones_replacer loaded.')
+                self.logger.log(logging.INFO, 'successfully loaded HomophonesReplacer.')
                 return True
             except (IOError, json.JSONDecodeError) as e:
-                self.logger.log(logging.WARNING, f'Error loading homophones map: {e}')
+                self.logger.log(logging.WARNING, f'error loading homophones map: {e}')
             except Exception as e:
-                self.logger.log(logging.WARNING, f'Error loading homophones_replacer: {e}')
+                self.logger.log(logging.WARNING, f'error loading HomophonesReplacer: {e}')
         return False
diff --git a/ChatTTS/utils/infer.py b/ChatTTS/utils/infer.py
index ddebd60..63fcf9b 100644
--- a/ChatTTS/utils/infer.py
+++ b/ChatTTS/utils/infer.py
@@ -1,8 +1,12 @@
-
+import json
 import re
+from typing import Dict, Tuple, List
+import sys
+
+from numba import jit
+import numpy as np
 import torch
 import torch.nn.functional as F
-import json
 
     
 class CustomRepetitionPenaltyLogitsProcessorRepeat():
@@ -47,6 +51,26 @@ class CustomRepetitionPenaltyLogitsProcessor():
         
         return scores
 
+@jit
+def _find_index(table: np.ndarray, val: np.uint16):
+    for i in range(table.size):
+        if table[i] == val:
+            return i
+    return -1
+
+@jit
+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):
+        ch = result[i]
+        p = _find_index(table[0], ch)
+        if p >= 0:
+            repl_char = table[1][p]
+            result[i] = repl_char
+            replaced_words.append((chr(ch), chr(repl_char)))
+    return result, replaced_words
+
 class HomophonesReplacer:
     """
     Homophones Replacer
@@ -65,37 +89,40 @@ class HomophonesReplacer:
     [python-pinyin](https://github.com/mozillazg/python-pinyin)
 
     """
-    def __init__(self, map_file_path):
-        self.homophones_map = self.load_homophones_map(map_file_path)
+    def __init__(self, map_file_path: str):
+        self.homophones_map = self._load_homophones_map(map_file_path)
+        self.coding = "utf-16-le" if sys.byteorder == "little" else "utf-16-be"
 
-    def load_homophones_map(self, map_file_path):
+    def _load_homophones_map(self, map_file_path: str) -> np.ndarray:
         with open(map_file_path, 'r', encoding='utf-8') as f:
-            homophones_map = json.load(f)
-        return homophones_map
+            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 replace(self, text):
-        result = []
-        replaced_words = []
-        for char in text:
-            if char in self.homophones_map:
-                repl_char = self.homophones_map[char]
-                result.append(repl_char)
-                replaced_words.append((char, repl_char))
-            else:
-                result.append(char)
-        return ''.join(result), replaced_words
+    def replace(self, text: str):
+        arr, lst = _fast_replace(
+            self.homophones_map,
+            text.encode(self.coding),
+        )
+        return arr.tobytes().decode(self.coding), lst
 
-def count_invalid_characters(s):
-    
-    s = re.sub(r'\[uv_break\]|\[laugh\]|\[lbreak\]', '', s)
-    pattern = re.compile(r'[^\u4e00-\u9fffA-Za-z，。、,\. ]')
-    non_alphabetic_chinese_chars = pattern.findall(s)
+accept_pattern = re.compile(r'[^\u4e00-\u9fffA-Za-z，。、,\. ]')
+sub_pattern = re.compile(r'\[uv_break\]|\[laugh\]|\[lbreak\]')
+
+def count_invalid_characters(s: str):
+    global accept_pattern, sub_pattern
+    s = sub_pattern.sub('', s)
+    non_alphabetic_chinese_chars = accept_pattern.findall(s)
     return set(non_alphabetic_chinese_chars)
 
-def detect_language(sentence):
+chinese_char_pattern = re.compile(r'[\u4e00-\u9fff]')
+english_word_pattern = re.compile(r'\b[A-Za-z]+\b')
 
-    chinese_char_pattern = re.compile(r'[\u4e00-\u9fff]')
-    english_word_pattern = re.compile(r'\b[A-Za-z]+\b')
+def detect_language(sentence):
+    global chinese_char_pattern, english_word_pattern
 
     chinese_chars = chinese_char_pattern.findall(sentence)
     english_words = english_word_pattern.findall(sentence)
@@ -104,9 +131,9 @@ def detect_language(sentence):
         return "zh"
     else:
         return "en"
-    
-    
-character_map = {
+
+
+character_simplifier = str.maketrans({
     '：': '，',
     '；': '，',
     '！': '。',
@@ -135,9 +162,9 @@ character_map = {
     '>': ',',
     '<': ',',
     '-': ',',
-}
+})
 
-halfwidth_2_fullwidth_map = {
+halfwidth_2_fullwidth = str.maketrans({
         '!': '！',
         '"': '“',
         "'": '‘',
@@ -170,12 +197,10 @@ halfwidth_2_fullwidth_map = {
         '|': '｜',
         '}': '｝',
         '~': '～'
-    }
+    })
 
-def apply_half2full_map(text):
-    translation_table = str.maketrans(halfwidth_2_fullwidth_map)
-    return text.translate(translation_table)
+def apply_half2full_map(text: str) -> str:
+    return text.translate(halfwidth_2_fullwidth)
 
-def apply_character_map(text):
-    translation_table = str.maketrans(character_map)
-    return text.translate(translation_table)
+def apply_character_map(text: str) -> str:
+    return text.translate(character_simplifier)
diff --git a/tools/audio/np.py b/tools/audio/np.py
index 3dfb46d..9812f78 100644
--- a/tools/audio/np.py
+++ b/tools/audio/np.py
@@ -7,8 +7,7 @@ def unsafe_float_to_int16(audio: np.ndarray) -> np.ndarray:
     This function will destroy audio, use only once.
     """
     am = np.abs(audio).max() * 32768
-    if am > 32768:
-        am = 32768 * 32768 / am
+    am = 32767 * 32768 / am
     np.multiply(audio, am, audio)
     audio16 = audio.astype(np.int16)
     return audio16

From 1b16feec2fc38293cd8378bdc06e3f3a63960e6f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Sun, 23 Jun 2024 22:45:20 +0900
Subject: [PATCH 7/9] chore(core): add more types

---
 ChatTTS/core.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/ChatTTS/core.py b/ChatTTS/core.py
index e3891cd..27f6085 100644
--- a/ChatTTS/core.py
+++ b/ChatTTS/core.py
@@ -3,8 +3,9 @@ import json
 import logging
 import tempfile
 from functools import partial
-from typing import Literal, Optional, List
+from typing import Literal, Optional, List, Callable
 
+import numpy as np
 import torch
 from omegaconf import OmegaConf
 from vocos import Vocos
@@ -134,11 +135,11 @@ class Chat:
             vocos.load_state_dict(torch.load(vocos_ckpt_path))
             self.vocos = vocos
             if "mps" in str(self.device):
-                self._vocos_decode = lambda spec: self.vocos.decode(
+                self._vocos_decode: Callable[[torch.Tensor], np.ndarray] = lambda spec: self.vocos.decode(
                     spec.cpu()
                 ).cpu().numpy()
             else:
-                self._vocos_decode = lambda spec: self.vocos.decode(
+                self._vocos_decode: Callable[[torch.Tensor], np.ndarray]  = lambda spec: self.vocos.decode(
                     spec
                 ).cpu().numpy()
             self.logger.log(logging.INFO, 'vocos loaded.')
@@ -295,7 +296,7 @@ class Chat:
 
     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 = []
+        wavs: List[np.ndarray] = []
         for i, chunk_data in enumerate(x):
             start_seek = start_seeks[i]
             length = len(chunk_data)

From d93bc19f414672cb0883b103f706a930f1dde1fe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Sun, 23 Jun 2024 23:37:38 +0900
Subject: [PATCH 8/9] optimize(dl): use mmap to hash

https://stackoverflow.com/a/67412295
---
 ChatTTS/utils/download.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/ChatTTS/utils/download.py b/ChatTTS/utils/download.py
index d26fae2..d9e4bf6 100644
--- a/ChatTTS/utils/download.py
+++ b/ChatTTS/utils/download.py
@@ -3,15 +3,15 @@ from pathlib import Path
 import hashlib
 import requests
 from io import BytesIO
+from mmap import mmap, ACCESS_READ
 
 from .log import logger
 
-def sha256(f) -> str:
-    sha256_hash = hashlib.sha256()
-    # Read and update hash in chunks of 4M
-    for byte_block in iter(lambda: f.read(4 * 1024 * 1024), b""):
-        sha256_hash.update(byte_block)
-    return sha256_hash.hexdigest()
+def sha256(fileno: int) -> str:
+    data = mmap(fileno, 0, access=ACCESS_READ)
+    h = hashlib.sha256(data).hexdigest()
+    del data
+    return h
 
 
 def check_model(
@@ -24,7 +24,7 @@ def check_model(
         logger.info(f"{target} not exist.")
         return False
     with open(target, "rb") as f:
-        digest = sha256(f)
+        digest = sha256(f.fileno())
         bakfile = f"{target}.bak"
         if digest != hash:
             logger.warn(f"{target} sha256 hash mismatch.")

From 51c2118b76a50843e3a32c3efbc32bc5ecb3e731 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Sun, 23 Jun 2024 23:45:50 +0900
Subject: [PATCH 9/9] fix(log): utils log cannot display

---
 ChatTTS/core.py                      |  6 ++---
 ChatTTS/utils/{download.py => dl.py} | 36 ++++++++++++++--------------
 ChatTTS/utils/gpu.py                 |  6 ++---
 ChatTTS/utils/io.py                  |  2 +-
 ChatTTS/utils/log.py                 | 14 +++++++----
 5 files changed, 35 insertions(+), 29 deletions(-)
 rename ChatTTS/utils/{download.py => dl.py} (81%)

diff --git a/ChatTTS/core.py b/ChatTTS/core.py
index 27f6085..a903ee7 100644
--- a/ChatTTS/core.py
+++ b/ChatTTS/core.py
@@ -17,8 +17,8 @@ from .utils.gpu import select_device
 from .utils.infer import count_invalid_characters, detect_language, apply_character_map, apply_half2full_map, HomophonesReplacer
 from .utils.io import get_latest_modified_file, del_all
 from .infer.api import refine_text, infer_code
-from .utils.download import check_all_assets, download_all_assets
-from .utils.log import set_utils_logger
+from .utils.dl import check_all_assets, download_all_assets
+from .utils.log import logger as utils_logger
 
 
 class Chat:
@@ -27,7 +27,7 @@ class Chat:
         self.normalizer = {}
         self.homophones_replacer = None
         self.logger = logger
-        set_utils_logger(logger)
+        utils_logger.set_logger(logger)
         
     def has_loaded(self, use_decoder = False):
         not_finish = False
diff --git a/ChatTTS/utils/download.py b/ChatTTS/utils/dl.py
similarity index 81%
rename from ChatTTS/utils/download.py
rename to ChatTTS/utils/dl.py
index d9e4bf6..d061f22 100644
--- a/ChatTTS/utils/download.py
+++ b/ChatTTS/utils/dl.py
@@ -19,18 +19,18 @@ def check_model(
 ) -> bool:
     target = dir_name / model_name
     relname = target.as_posix()
-    logger.debug(f"checking {relname}...")
+    logger.get_logger().debug(f"checking {relname}...")
     if not os.path.exists(target):
-        logger.info(f"{target} not exist.")
+        logger.get_logger().info(f"{target} not exist.")
         return False
     with open(target, "rb") as f:
         digest = sha256(f.fileno())
         bakfile = f"{target}.bak"
         if digest != hash:
-            logger.warn(f"{target} sha256 hash mismatch.")
-            logger.info(f"expected: {hash}")
-            logger.info(f"real val: {digest}")
-            logger.warn("please add parameter --update to download the latest assets.")
+            logger.get_logger().warn(f"{target} sha256 hash mismatch.")
+            logger.get_logger().info(f"expected: {hash}")
+            logger.get_logger().info(f"real val: {digest}")
+            logger.get_logger().warn("please add parameter --update to download the latest assets.")
             if remove_incorrect:
                 if not os.path.exists(bakfile):
                     os.rename(str(target), bakfile)
@@ -45,7 +45,7 @@ def check_model(
 def check_all_assets(update=False) -> bool:
     BASE_DIR = Path(os.getcwd())
 
-    logger.info("checking assets...")
+    logger.get_logger().info("checking assets...")
     current_dir = BASE_DIR / "asset"
     names = [
         "Decoder.pt",
@@ -62,7 +62,7 @@ def check_all_assets(update=False) -> bool:
         ):
             return False
 
-    logger.info("checking configs...")
+    logger.get_logger().info("checking configs...")
     current_dir = BASE_DIR / "config"
     names = [
         "decoder.yaml",
@@ -78,44 +78,44 @@ def check_all_assets(update=False) -> bool:
         ):
             return False
 
-    logger.info("all assets are already latest.")
+    logger.get_logger().info("all assets are already latest.")
     return True
 
 
 def download_and_extract_tar_gz(url: str, folder: str):
     import tarfile
 
-    logger.info(f"downloading {url}")
+    logger.get_logger().info(f"downloading {url}")
     response = requests.get(url, stream=True, timeout=(5, 10))
     with BytesIO() as out_file:
         out_file.write(response.content)
         out_file.seek(0)
-        logger.info(f"downloaded.")
+        logger.get_logger().info(f"downloaded.")
         with tarfile.open(fileobj=out_file, mode="r:gz") as tar:
             tar.extractall(folder)
-        logger.info(f"extracted into {folder}")
+        logger.get_logger().info(f"extracted into {folder}")
 
 
 def download_and_extract_zip(url: str, folder: str):
     import zipfile
 
-    logger.info(f"downloading {url}")
+    logger.get_logger().info(f"downloading {url}")
     response = requests.get(url, stream=True, timeout=(5, 10))
     with BytesIO() as out_file:
         out_file.write(response.content)
         out_file.seek(0)
-        logger.info(f"downloaded.")
+        logger.get_logger().info(f"downloaded.")
         with zipfile.ZipFile(out_file) as zip_ref:
             zip_ref.extractall(folder)
-        logger.info(f"extracted into {folder}")
+        logger.get_logger().info(f"extracted into {folder}")
 
 
 def download_dns_yaml(url: str, folder: str):
-    logger.info(f"downloading {url}")
+    logger.get_logger().info(f"downloading {url}")
     response = requests.get(url, stream=True, timeout=(5, 10))
     with open(os.path.join(folder, "dns.yaml"), "wb") as out_file:
         out_file.write(response.content)
-        logger.info(f"downloaded into {folder}")
+        logger.get_logger().info(f"downloaded into {folder}")
 
 
 def download_all_assets(tmpdir: str, version="0.2.5"):
@@ -140,7 +140,7 @@ def download_all_assets(tmpdir: str, version="0.2.5"):
 
     architecture = archs.get(architecture, None)
     if not architecture:
-        logger.error(f"architecture {architecture} is not supported")
+        logger.get_logger().error(f"architecture {architecture} is not supported")
         exit(1)
     try:
         BASE_URL = "https://github.com/fumiama/RVC-Models-Downloader/releases/download/"
diff --git a/ChatTTS/utils/gpu.py b/ChatTTS/utils/gpu.py
index 8c39694..9d3836e 100644
--- a/ChatTTS/utils/gpu.py
+++ b/ChatTTS/utils/gpu.py
@@ -14,14 +14,14 @@ def select_device(min_memory=2048):
         device = torch.device(f'cuda:{selected_gpu}')
         free_memory_mb = max_free_memory / (1024 * 1024)
         if free_memory_mb < min_memory:
-            logger.warning(f'GPU {selected_gpu} has {round(free_memory_mb, 2)} MB memory left. Switching to 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.info('Apple GPU found, using MPS.')
+        logger.get_logger().info('Apple GPU found, using MPS.')
         device = torch.device('mps')
     else:
-        logger.warning('No GPU found, use CPU instead')
+        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 eeb9871..b787d21 100644
--- a/ChatTTS/utils/io.py
+++ b/ChatTTS/utils/io.py
@@ -10,7 +10,7 @@ def get_latest_modified_file(directory):
 
     files = [os.path.join(directory, f) for f in os.listdir(directory)] 
     if not files:
-        logger.log(logging.WARNING, f'no files found in the directory: {directory}')
+        logger.get_logger().log(logging.WARNING, f'no files found in the directory: {directory}')
         return None
     latest_file = max(files, key=os.path.getmtime)
 
diff --git a/ChatTTS/utils/log.py b/ChatTTS/utils/log.py
index 6055f84..382d79f 100644
--- a/ChatTTS/utils/log.py
+++ b/ChatTTS/utils/log.py
@@ -1,8 +1,14 @@
 import logging
 from pathlib import Path
 
-logger = logging.getLogger(Path(__file__).parent.name)
+class Logger():
+    def __init__(self, logger=logging.getLogger(Path(__file__).parent.name)):
+        self.logger = logger
 
-def set_utils_logger(l: logging.Logger):
-    global logger
-    logger = l
+    def set_logger(self, logger: logging.Logger):
+        self.logger = logger
+    
+    def get_logger(self) -> logging.Logger:
+        return self.logger
+
+logger = Logger()