拆分gpu、cpu模式运行模式

Chat2TTS/core.py

@@ -1,4 +1,3 @@
-
 import os
 import logging
 from omegaconf import OmegaConf
@@ -11,9 +10,11 @@ from .utils.gpu_utils import select_device
 from .utils.io_utils import get_latest_modified_file
 from .infer.api import refine_text, infer_code
 from dataclasses import dataclass
-from typing import Literal, Optional, List, Tuple, Dict
+from typing import Literal, Optional, List, Tuple, Dict, Union
+import numpy as np
 from tool.logger import get_logger
-from tool.normalizer import normalizer_en_nemo_text,normalizer_cn_tn
+from tool.normalizer import normalizer_en_nemo_text, normalizer_cn_tn
+from tool.func import encode_prompt
 
 from ChatTTS.norm import Normalizer
 
@@ -23,31 +24,31 @@ from huggingface_hub import snapshot_download
 class Chat:
     def __init__(self, ):
         self.pretrain_models = {}
-        self.logger = get_logger(__name__,lv=logging.INFO)
+        self.logger = get_logger(__name__, lv=logging.INFO)
         self.normalizer = Normalizer(
             os.path.join(os.path.dirname(__file__), "res", "homophones_map.json"),
             self.logger,
         )
-        
-    def check_model(self, level = logging.INFO, use_decoder = False):
+
+    def check_model(self, level=logging.INFO, use_decoder=False):
         not_finish = False
         check_list = ['vocos', 'gpt', 'tokenizer']
-        
+
         if use_decoder:
             check_list.append('decoder')
         else:
             check_list.append('dvae')
-            
+
         for module in check_list:
             if module not in self.pretrain_models:
                 self.logger.log(logging.WARNING, f'{module} not initialized.')
                 not_finish = True
-                
+
         if not not_finish:
             self.logger.log(level, f'All initialized.')
-            
+
         return not not_finish
-        
+
     def load_models(self, source='huggingface', force_redownload=False, local_path='<LOCAL_PATH>'):
         if source == 'huggingface':
             hf_home = os.getenv('HF_HOME', os.path.expanduser("~/.cache/huggingface"))
@@ -55,25 +56,27 @@ class Chat:
                 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: 
+            if download_path is None or force_redownload:
                 self.logger.log(logging.INFO, f'Download from HF: https://huggingface.co/2Noise/ChatTTS')
                 download_path = snapshot_download(repo_id="2Noise/ChatTTS", allow_patterns=["*.pt", "*.yaml"])
             else:
                 self.logger.log(logging.INFO, f'Load from cache: {download_path}')
-            self._load(**{k: os.path.join(download_path, v) for k, v in OmegaConf.load(os.path.join(download_path, 'config', 'path.yaml')).items()})
+            self._load(**{k: os.path.join(download_path, v) for k, v in
+                          OmegaConf.load(os.path.join(download_path, 'config', 'path.yaml')).items()})
             self._regist_normalizer()
         elif source == 'local':
             self.logger.log(logging.INFO, f'Load from local: {local_path}')
-            self._load(**{k: os.path.join(local_path, v) for k, v in OmegaConf.load(os.path.join(local_path, 'config', 'path.yaml')).items()})
+            self._load(**{k: os.path.join(local_path, v) for k, v in
+                          OmegaConf.load(os.path.join(local_path, 'config', 'path.yaml')).items()})
 
     def _regist_normalizer(self):
 
         self.logger.info("==========开始注册 normalizer===========")
 
         try:
-             self.normalizer.register("en",normalizer_en_nemo_text())
+            self.normalizer.register("en", normalizer_en_nemo_text())
         except ValueError as e:
-            self.logger.error('normalizer_en_nemo_text register fail' , e)
+            self.logger.error('normalizer_en_nemo_text register fail', e)
         except:
             self.logger.error("Package nemo_text_processing not found!")
             self.logger.error(
@@ -81,40 +84,40 @@ class Chat:
             )
 
         try:
-            self.normalizer.register("zh",normalizer_cn_tn())
+            self.normalizer.register("zh", normalizer_cn_tn())
         except ValueError as e:
-            self.logger.error('normalizer_cn_tn register fail' , e)
+            self.logger.error('normalizer_cn_tn register fail', e)
         except:
             self.logger.error("Package WeTextProcessing not found!")
             self.logger.error(
                 "Run: conda install -c conda-forge pynini=2.1.5 && pip install WeTextProcessing",
             )
 
-
     def _load(
-        self, 
-        vocos_config_path: str = None, 
-        vocos_ckpt_path: str = None,
-        dvae_config_path: str = None,
-        dvae_ckpt_path: str = None,
-        gpt_config_path: str = None,
-        gpt_ckpt_path: str = None,
-        decoder_config_path: str = None,
-        decoder_ckpt_path: str = None,
-        tokenizer_path: str = None,
-        device: str = None
+            self,
+            vocos_config_path: str = None,
+            vocos_ckpt_path: str = None,
+            dvae_config_path: str = None,
+            dvae_ckpt_path: str = None,
+            gpt_config_path: str = None,
+            gpt_ckpt_path: str = None,
+            decoder_config_path: str = None,
+            decoder_ckpt_path: str = None,
+            tokenizer_path: str = None,
+            device: str = None
     ):
         if not device:
             device = select_device(4096)
             self.logger.log(logging.INFO, f'use {device}')
-            
+
+        self.device = device
         if vocos_config_path:
             vocos = Vocos.from_hparams(vocos_config_path).to(device).eval()
             assert vocos_ckpt_path, 'vocos_ckpt_path should not be None'
             vocos.load_state_dict(torch.load(vocos_ckpt_path))
             self.pretrain_models['vocos'] = vocos
             self.logger.log(logging.INFO, 'vocos loaded.')
-        
+
         if dvae_config_path:
             cfg = OmegaConf.load(dvae_config_path)
             dvae = DVAE(**cfg).to(device).eval()
@@ -122,7 +125,7 @@ class Chat:
             dvae.load_state_dict(torch.load(dvae_ckpt_path, map_location='cpu'))
             self.pretrain_models['dvae'] = dvae
             self.logger.log(logging.INFO, 'dvae loaded.')
-            
+
         if gpt_config_path:
             cfg = OmegaConf.load(gpt_config_path)
             gpt = GPT_warpper(**cfg).to(device).eval()
@@ -139,7 +142,6 @@ class Chat:
                 spk_stat_path, weights_only=True, mmap=True, map_location='cpu'
             ).to(device)
 
-
         if decoder_config_path:
             cfg = OmegaConf.load(decoder_config_path)
             decoder = DVAE(**cfg).to(device).eval()
@@ -147,13 +149,13 @@ class Chat:
             decoder.load_state_dict(torch.load(decoder_ckpt_path, map_location='cpu'))
             self.pretrain_models['decoder'] = decoder
             self.logger.log(logging.INFO, 'decoder loaded.')
-        
+
         if tokenizer_path:
             tokenizer = torch.load(tokenizer_path, map_location='cpu')
             tokenizer.padding_side = 'left'
             self.pretrain_models['tokenizer'] = tokenizer
             self.logger.log(logging.INFO, 'tokenizer loaded.')
-            
+
         self.check_model()
 
     @dataclass(repr=False, eq=False)
@@ -177,16 +179,19 @@ class Chat:
         max_new_token: int = 2048
 
     def infer(
-        self, 
-        text, 
-        skip_refine_text=False, 
-        refine_text_only=False, 
-        params_refine_text={}, 
-        params_infer_code={}, 
-        use_decoder=False,
-        lang=None
+            self,
+            text,
+            skip_refine_text=False,
+            refine_text_only=False,
+            params_refine_text={},
+            params_infer_code={},
+            use_decoder=False,
+            lang=None
     ):
-        
+
+        self.logger.info(
+            f"========开始infer模型,use_decoder:{use_decoder},lang:{lang},"
+            f"mskip_refine_text:{skip_refine_text},refine_text_only:{refine_text_only}======")
         assert self.check_model(use_decoder=use_decoder)
 
         if not isinstance(text, list):
@@ -203,36 +208,48 @@ class Chat:
         ]
 
         if skip_refine_text:
-            self.logger.info(f"========对文本内容不做优化处理,仅做规则处理,lang:{lang}======")
+            self.logger.info(f"========对文本内容不做优化处理,仅做规则处理======")
         else:
             self.logger.info(f"========针对文本内容做模型优化处理,lang:{lang}======")
             text_tokens = refine_text(self.pretrain_models, text, **params_refine_text)['ids']
-            text_tokens = [i[i < self.pretrain_models['tokenizer'].convert_tokens_to_ids('[break_0]')] for i in text_tokens]
+            text_tokens = [i[i < self.pretrain_models['tokenizer'].convert_tokens_to_ids('[break_0]')] for i in
+                           text_tokens]
             text = self.pretrain_models['tokenizer'].batch_decode(text_tokens)
             if refine_text_only:
                 return text
-            
+
         text = [params_infer_code.get('prompt', '') + i for i in text]
         params_infer_code.pop('prompt', '')
         result = infer_code(self.pretrain_models, text, **params_infer_code, return_hidden=use_decoder)
-        
+
         if use_decoder:
-            mel_spec = [self.pretrain_models['decoder'](i[None].permute(0,2,1)) for i in result['hiddens']]
+            mel_spec = [self.pretrain_models['decoder'](i[None].permute(0, 2, 1)) for i in result['hiddens']]
         else:
-            mel_spec = [self.pretrain_models['dvae'](i[None].permute(0,2,1)) for i in result['ids']]
-            
+            mel_spec = [self.pretrain_models['dvae'](i[None].permute(0, 2, 1)) for i in result['ids']]
+
         wav = [self.pretrain_models['vocos'].decode(i).cpu().numpy() for i in mel_spec]
-        
+
         return wav
 
+    # 返回一个空的wav 音频文件
     def emptpy_audio(self):
-        return  self.infer(" ",
-                skip_refine_text=True,
-                refine_text_only=False,
-                params_refine_text={},
-                params_infer_code={},
-                use_decoder=False)
-        
+        return self.infer(" ",
+                          skip_refine_text=True,
+                          refine_text_only=False,
+                          params_refine_text={},
+                          params_infer_code={},
+                          use_decoder=False)
+
+    '''
+     将音频张量 做转码处理
+    '''
+
+    @torch.inference_mode()
+    def sample_audio_speaker(self, wav: Union[np.ndarray, torch.Tensor]) -> str:
+        if isinstance(wav, np.ndarray):
+            wav = torch.from_numpy(wav).to(self.device)
+        squeeze = self.pretrain_models['dvae'](wav, "encode").squeeze_(0)
+        return encode_prompt(squeeze)
 
     # def sample_random_speaker(self) -> str:
     #     return self._encode_spk_emb(self.sample_random_speaker_tensor())
@@ -266,4 +283,4 @@ class Chat:
                 .add_(mean)
             )
             del out, std, mean
-            return spk
+            return spk

test/audio_test.py

@@ -0,0 +1,48 @@
+import os, sys
+
+if sys.platform == "darwin":
+    os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+now_dir = os.getcwd()
+sys.path.append(now_dir)
+
+import Chat2TTS
+from  tool.av import load_audio
+from tool.logger import get_logger
+
+
+
+logger = get_logger("audio_test")
+# Initialize and load the model:
+chat = Chat2TTS.Chat()
+
+def init_chat():
+    global chat
+    source = "local"
+    # 获取启动模式
+    MODEL = os.getenv('MODEL')
+    # huggingface 部署模式下，模型则直接使用hf的模型数据
+    if MODEL == "HF":
+        source = "huggingface"
+
+    logger.info("loading Chat2TTS model..., start source:" + source)
+
+
+    if chat.load_models(source=source, local_path="D:\\chenjgspace\\ai-model\\chattts"):
+        print("Models loaded successfully.")
+    logger.info("Models loaded end.")
+    # else:
+    #     logger.error("=========Models load failed.")
+    #     sys.exit(1)
+
+def audo_encode():
+    sample_audio = load_audio("D:\\Download\\audio_test.wav",24000)
+    logger.info("================sample_audio："+str(sample_audio))
+    spk_smp=chat.sample_audio_speaker(sample_audio)
+    logger.info("================spk_smp："+str(spk_smp))
+
+
+if __name__ == "__main__":
+
+    init_chat()
+    # 还需要继续调试
+    audo_encode()

test/common_test.py

@@ -8,7 +8,7 @@ from tool.logger import get_logger
 
 logger=get_logger("common-test")
 def save_mp3_file(wav, index, prefix_name):
-    from tool.pcm import pcm_arr_to_mp3_view
+    from tool.np import pcm_arr_to_mp3_view
     data = pcm_arr_to_mp3_view(wav)
     mp3_filename = prefix_name + "_" + str(index) + ".mp3"
     with open(mp3_filename, "wb") as f:

tool/__init__.py

@@ -1,5 +1,4 @@
 from .av import load_audio
-from .pcm import pcm_arr_to_mp3_view
-from .np import float_to_int16
+from .np import float_to_int16,pcm_arr_to_mp3_view
 from .ctx import TorchSeedContext
 from .gpu import select_device

tool/func.py

@@ -1,6 +1,11 @@
 
 import gradio as gr
 import random
+import torch
+import lzma
+
+import numpy as np
+import pybase16384 as b14
 
 seed_min = 1
 seed_max = 4294967295
@@ -30,6 +35,28 @@ voices = {
 def on_voice_change(vocie_selection):
     return voices.get(vocie_selection)["seed"]
 
-
+'''
+随机生成种子
+'''
 def generate_seed():
-    return gr.update(value=random.randint(seed_min, seed_max))
+    return gr.update(value=random.randint(seed_min, seed_max))
+
+'''
+ 音频文件张量 编码
+'''
+
+@torch.no_grad()
+def encode_prompt(prompt: torch.Tensor) -> str:
+    arr: np.ndarray = prompt.to(dtype=torch.uint16, device="cpu").numpy()
+    shp = arr.shape
+    assert len(shp) == 2, "prompt must be a 2D tensor"
+    s = b14.encode_to_string(
+        np.array(shp, dtype="<u2").tobytes()
+        + lzma.compress(
+            arr.astype("<u2").tobytes(),
+            format=lzma.FORMAT_RAW,
+            filters=[{"id": lzma.FILTER_LZMA2, "preset": 9 | lzma.PRESET_EXTREME}],
+        ),
+    )
+    del arr
+    return s

tool/np.py

@@ -1,11 +1,28 @@
 import math
 
-import numpy as np
 from numba import jit
+import wave
+from io import BytesIO
+
+import numpy as np
+from .av import wav2
+
 
 
-@jit
 def float_to_int16(audio: np.ndarray) -> np.ndarray:
     am = int(math.ceil(float(np.abs(audio).max())) * 32768)
     am = 32767 * 32768 // am
     return np.multiply(audio, am).astype(np.int16)
+
+def pcm_arr_to_mp3_view(wav: np.ndarray):
+    buf = BytesIO()
+    with wave.open(buf, "wb") as wf:
+        wf.setnchannels(1)  # Mono channel
+        wf.setsampwidth(2)  # Sample width in bytes
+        wf.setframerate(24000)  # Sample rate in Hz
+        wf.writeframes(float_to_int16(wav))
+    buf.seek(0, 0)
+    buf2 = BytesIO()
+    wav2(buf, buf2, "mp3")
+    buf.seek(0, 0)
+    return buf2.getbuffer()

tool/pcm.py

@@ -1,21 +0,0 @@
-import wave
-from io import BytesIO
-
-import numpy as np
-
-from .np import float_to_int16
-from .av import wav2
-
-
-def pcm_arr_to_mp3_view(wav: np.ndarray):
-    buf = BytesIO()
-    with wave.open(buf, "wb") as wf:
-        wf.setnchannels(1)  # Mono channel
-        wf.setsampwidth(2)  # Sample width in bytes
-        wf.setframerate(24000)  # Sample rate in Hz
-        wf.writeframes(float_to_int16(wav))
-    buf.seek(0, 0)
-    buf2 = BytesIO()
-    wav2(buf, buf2, "mp3")
-    buf.seek(0, 0)
-    return buf2.getbuffer()

web/app_cpu.py

@@ -45,7 +45,7 @@ def init_chat(args):
 
 def main(args):
     with gr.Blocks() as demo:
-        gr.Markdown("# ChatTTS demo")
+        gr.Markdown("# ChatTTS demo CPU模式下运行")
         with gr.Row():
             with gr.Column(scale=1):
                 text_input = gr.Textbox(

web/app_gpu.py

@@ -48,7 +48,7 @@ def init_chat(args):
 
 def main(args):
     with gr.Blocks() as demo:
-        gr.Markdown("# ChatTTS demo")
+        gr.Markdown("# ChatTTS demo GPU模式下运行")
         with gr.Row():
             with gr.Column(scale=1):
                 text_input = gr.Textbox(
@@ -71,6 +71,12 @@ def main(args):
                     interactive=True,
                     value=True
                 )
+
+            use_decoder_checkBox = gr.Checkbox(
+                label="是否使用decoder模型,如否则使用dvae模型",
+                interactive=True,
+                value=True
+            )
             temperature_slider = gr.Slider(
                 minimum=0.00001,
                 maximum=1.0,
@@ -79,22 +85,23 @@ def main(args):
                 interactive=True,
                 label="模型 Temperature 参数设置"
             )
-            top_p_slider = gr.Slider(
-                minimum=0.1,
-                maximum=0.9,
-                step=0.05,
-                value=0.7,
-                label="模型 top_P 参数设置",
-                interactive=True,
-            )
-            top_k_slider = gr.Slider(
-                minimum=1,
-                maximum=20,
-                step=1,
-                value=20,
-                label="模型 top_K 参数设置",
-                interactive=True,
-            )
+            with gr.Column():
+                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():
             lang_selection = gr.Dropdown(
                 label="语种",
@@ -139,7 +146,7 @@ def main(args):
             # )
 
         with gr.Row():
-            reload_chat_button = gr.Button("Reload", scale=1, interactive=True)
+           # reload_chat_button = gr.Button("Reload", scale=1, interactive=True)
             generate_button = gr.Button("生成音频文件", scale=1, interactive=True)
 
         with gr.Row():
@@ -177,11 +184,12 @@ def main(args):
                                                        text_seed_input,
                                                        refine_text_checkBox,
                                                        refine_audio_checkBox,
+                                                       use_decoder_checkBox,
                                                        temperature_slider,
                                                        top_p_slider,
                                                        top_k_slider,
                                                        audio_seed_input,
-                                                        lang_selection
+                                                       lang_selection
                                                        ],
                               outputs=[text_output,audio_output])
         # 初始化 spk_emb_text 数值
@@ -212,6 +220,7 @@ def general_chat_infer_audio(text,
                              text_seed_input,
                              refine_text_checkBox,
                              refine_audio_checkBox,
+                             use_decoder_checkBox,
                              temperature_slider,
                              top_p_slider,
                              top_k_slider,
@@ -239,7 +248,8 @@ def general_chat_infer_audio(text,
                 skip_refine_text=False,
                 refine_text_only=True,  #仅返回优化后文本内容
                 params_refine_text=params_refine_text,
-                lang=lang
+                lang=lang,
+                use_decoder=use_decoder_checkBox
             )
 
 
@@ -265,6 +275,7 @@ def general_chat_infer_audio(text,
                 skip_refine_text=True, #跳过文本优化
                 params_refine_text=params_refine_text,
                 params_infer_code=params_infer_code,
+                use_decoder=use_decoder_checkBox
             )
 
     #yield 24000, float_to_int16(wav[0]).T