From 5f21aef7860a3a7af140bc93f294e155ac7b5bcb Mon Sep 17 00:00:00 2001 From: "zhoubofan.zbf" Date: Thu, 29 Aug 2024 23:35:07 +0800 Subject: [PATCH 1/5] add flow decoder tensorrt infer --- cosyvoice/bin/export_trt.py | 105 ++++++++++++++++++++++++++++++-- cosyvoice/cli/cosyvoice.py | 5 +- cosyvoice/cli/model.py | 23 +++++-- cosyvoice/flow/decoder.py | 2 +- cosyvoice/flow/flow_matching.py | 33 ++++++++-- 5 files changed, 149 insertions(+), 19 deletions(-) diff --git a/cosyvoice/bin/export_trt.py b/cosyvoice/bin/export_trt.py index e6d480c..fea8205 100644 --- a/cosyvoice/bin/export_trt.py +++ b/cosyvoice/bin/export_trt.py @@ -1,8 +1,103 @@ -# TODO 跟export_jit一样的逻辑，完成flow部分的estimator的onnx导出。 -# tensorrt的安装方式，再这里写一下步骤提示如下，如果没有安装，那么不要执行这个脚本，提示用户先安装，不给选择 +import argparse +import logging +import os +import sys + +logging.getLogger('matplotlib').setLevel(logging.WARNING) + try: import tensorrt except ImportError: - print('step1, 下载\n step2. 解压，安装whl，') -# 安装命令里tensosrt的根目录用环境变量导入，比如os.environ['tensorrt_root_dir']/bin/exetrace，然后python里subprocess里执行导出命令 -# 后面我会在run.sh里写好执行命令 tensorrt_root_dir=xxxx python cosyvoice/bin/export_trt.py --model_dir xxx \ No newline at end of file + error_msg_zh = [ + "step.1 下载 tensorrt .tar.gz 压缩包并解压，下载地址: https://developer.nvidia.com/tensorrt/download/10x", + "step.2 使用 tensorrt whl 包进行安装根据 python 版本对应进行安装，如 pip install ${TensorRT-Path}/python/tensorrt-10.2.0-cp38-none-linux_x86_64.whl", + "step.3 将 tensorrt 的 lib 路径添加进环境变量中，export LD_LIBRARY_PATH=${TensorRT-Path}/lib/" + ] + print("\n".join(error_msg_zh)) + sys.exit(1) + +import torch +from cosyvoice.cli.cosyvoice import CosyVoice + +def get_args(): + parser = argparse.ArgumentParser(description='Export your model for deployment') + parser.add_argument('--model_dir', + type=str, + default='pretrained_models/CosyVoice-300M', + help='Local path to the model directory') + + parser.add_argument('--export_half', + action='store_true', + help='Export with half precision (FP16)') + + args = parser.parse_args() + print(args) + return args + +def main(): + args = get_args() + + cosyvoice = CosyVoice(args.model_dir, load_jit=False, load_trt=False) + + flow = cosyvoice.model.flow + estimator = cosyvoice.model.flow.decoder.estimator + + dtype = torch.float32 if not args.export_half else torch.float16 + device = torch.device("cuda") + batch_size = 1 + seq_len = 1024 + hidden_size = flow.output_size + x = torch.rand((batch_size, hidden_size, seq_len), dtype=dtype, device=device) + mask = torch.zeros((batch_size, 1, seq_len), dtype=dtype, device=device) + mu = torch.rand((batch_size, hidden_size, seq_len), dtype=dtype, device=device) + t = torch.tensor([0.], dtype=dtype, device=device) + spks = torch.rand((batch_size, hidden_size), dtype=dtype, device=device) + cond = torch.rand((batch_size, hidden_size, seq_len), dtype=dtype, device=device) + + onnx_file_name = 'estimator_fp16.onnx' if args.export_half else 'estimator_fp32.onnx' + onnx_file_path = os.path.join(args.model_dir, onnx_file_name) + dummy_input = (x, mask, mu, t, spks, cond) + + estimator = estimator.to(dtype) + + torch.onnx.export( + estimator, + dummy_input, + onnx_file_path, + export_params=True, + opset_version=18, + do_constant_folding=True, + input_names=['x', 'mask', 'mu', 't', 'spks', 'cond'], + output_names=['output'], + dynamic_axes={ + 'x': {2: 'seq_len'}, + 'mask': {2: 'seq_len'}, + 'mu': {2: 'seq_len'}, + 'cond': {2: 'seq_len'}, + 'output': {2: 'seq_len'}, + } + ) + + tensorrt_path = os.environ.get('tensorrt_root_dir') + if not tensorrt_path: + raise EnvironmentError("Please set the 'tensorrt_root_dir' environment variable.") + + if not os.path.isdir(tensorrt_path): + raise FileNotFoundError(f"The directory {tensorrt_path} does not exist.") + + trt_lib_path = os.path.join(tensorrt_path, "lib") + if trt_lib_path not in os.environ.get('LD_LIBRARY_PATH', ''): + print(f"Adding TensorRT lib path {trt_lib_path} to LD_LIBRARY_PATH.") + os.environ['LD_LIBRARY_PATH'] = f"{os.environ.get('LD_LIBRARY_PATH', '')}:{trt_lib_path}" + + trt_file_name = 'estimator_fp16.plan' if args.export_half else 'estimator_fp32.plan' + trt_file_path = os.path.join(args.model_dir, trt_file_name) + + trtexec_cmd = f"{tensorrt_path}/bin/trtexec --onnx={onnx_file_path} --saveEngine={trt_file_path} " \ + "--minShapes=x:1x80x1,mask:1x1x1,mu:1x80x1,t:1,spks:1x80,cond:1x80x1 " \ + "--maxShapes=x:1x80x4096,mask:1x1x4096,mu:1x80x4096,t:1,spks:1x80,cond:1x80x4096 --verbose" + + os.system(trtexec_cmd) + +if __name__ == "__main__": + main() diff --git a/cosyvoice/cli/cosyvoice.py b/cosyvoice/cli/cosyvoice.py index d5fbd4e..eac8922 100644 --- a/cosyvoice/cli/cosyvoice.py +++ b/cosyvoice/cli/cosyvoice.py @@ -21,7 +21,7 @@ from cosyvoice.utils.file_utils import logging class CosyVoice: - def __init__(self, model_dir, load_jit=True, load_trt=True): + def __init__(self, model_dir, load_jit=True, load_trt=True, use_fp16=False): instruct = True if '-Instruct' in model_dir else False self.model_dir = model_dir if not os.path.exists(model_dir): @@ -43,8 +43,7 @@ class CosyVoice: self.model.load_jit('{}/llm.text_encoder.fp16.zip'.format(model_dir), '{}/llm.llm.fp16.zip'.format(model_dir)) if load_trt: - # TODO - self.model.load_trt() + self.model.load_trt(model_dir, use_fp16) del configs def list_avaliable_spks(self): diff --git a/cosyvoice/cli/model.py b/cosyvoice/cli/model.py index 1184f0d..f074fd8 100644 --- a/cosyvoice/cli/model.py +++ b/cosyvoice/cli/model.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import os import torch import numpy as np import threading @@ -19,6 +20,10 @@ from contextlib import nullcontext import uuid from cosyvoice.utils.common import fade_in_out +try: + import tensorrt as trt +except ImportError: + ... class CosyVoiceModel: @@ -66,10 +71,20 @@ class CosyVoiceModel: llm_llm = torch.jit.load(llm_llm_model) self.llm.llm = llm_llm - def load_trt(self): - # TODO 你需要的TRT推理的准备 - self.flow.decoder.estimator = xxx - self.flow.decoder.session = xxx + def load_trt(self, model_dir, use_fp16): + trt_file_name = 'estimator_fp16.plan' if use_fp16 else 'estimator_fp32.plan' + trt_file_path = os.path.join(model_dir, trt_file_name) + if not os.path.isfile(trt_file_path): + raise f"{trt_file_path} does not exist. Please use bin/export_trt.py to generate .plan file" + + trt.init_libnvinfer_plugins(None, "") + logger = trt.Logger(trt.Logger.WARNING) + runtime = trt.Runtime(logger) + with open(trt_file_path, 'rb') as f: + serialized_engine = f.read() + engine = runtime.deserialize_cuda_engine(serialized_engine) + self.flow.decoder.estimator_context = engine.create_execution_context() + self.flow.decoder.estimator_engine = engine def llm_job(self, text, prompt_text, llm_prompt_speech_token, llm_embedding, uuid): with self.llm_context: diff --git a/cosyvoice/flow/decoder.py b/cosyvoice/flow/decoder.py index 4349279..be063d3 100755 --- a/cosyvoice/flow/decoder.py +++ b/cosyvoice/flow/decoder.py @@ -159,7 +159,7 @@ class ConditionalDecoder(nn.Module): _type_: _description_ """ - t = self.time_embeddings(t) + t = self.time_embeddings(t).to(t.dtype) t = self.time_mlp(t) x = pack([x, mu], "b * t")[0] diff --git a/cosyvoice/flow/flow_matching.py b/cosyvoice/flow/flow_matching.py index bcbaeb5..8cab545 100755 --- a/cosyvoice/flow/flow_matching.py +++ b/cosyvoice/flow/flow_matching.py @@ -30,6 +30,9 @@ class ConditionalCFM(BASECFM): # Just change the architecture of the estimator here self.estimator = estimator + self.estimator_context = None + self.estimator_engine = None + @torch.inference_mode() def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None): """Forward diffusion @@ -50,7 +53,7 @@ class ConditionalCFM(BASECFM): shape: (batch_size, n_feats, mel_timesteps) """ z = torch.randn_like(mu) * temperature - t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device) + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) if self.t_scheduler == 'cosine': t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond) @@ -71,6 +74,7 @@ class ConditionalCFM(BASECFM): cond: Not used but kept for future purposes """ t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0] + t = t.unsqueeze(dim=0) # I am storing this because I can later plot it by putting a debugger here and saving it to a file # Or in future might add like a return_all_steps flag @@ -96,13 +100,30 @@ class ConditionalCFM(BASECFM): return sol[-1] - # TODO - def forward_estimator(self): - if isinstance(self.estimator, trt): + def forward_estimator(self, x, mask, mu, t, spks, cond): + if self.estimator_context is not None: assert self.training is False, 'tensorrt cannot be used in training' - return xxx + bs = x.shape[0] + hs = x.shape[1] + seq_len = x.shape[2] + # assert bs == 1 and hs == 80 + ret = torch.empty_like(x) + self.estimator_context.set_input_shape("x", x.shape) + self.estimator_context.set_input_shape("mask", mask.shape) + self.estimator_context.set_input_shape("mu", mu.shape) + self.estimator_context.set_input_shape("t", t.shape) + self.estimator_context.set_input_shape("spks", spks.shape) + self.estimator_context.set_input_shape("cond", cond.shape) + bindings = [x.data_ptr(), mask.data_ptr(), mu.data_ptr(), t.data_ptr(), spks.data_ptr(), cond.data_ptr(), ret.data_ptr()] + + for i in range(len(bindings)): + self.estimator_context.set_tensor_address(self.estimator_engine.get_tensor_name(i), bindings[i]) + + handle = torch.cuda.current_stream().cuda_stream + self.estimator_context.execute_async_v3(stream_handle=handle) + return ret else: - return self.estimator.forward + return self.estimator.forward(x, mask, mu, t, spks, cond) def compute_loss(self, x1, mask, mu, spks=None, cond=None): """Computes diffusion loss From 53a3c1b17fda78da56be12ac9bec3739467f90db Mon Sep 17 00:00:00 2001 From: "zhoubofan.zbf" Date: Fri, 30 Aug 2024 00:47:40 +0800 Subject: [PATCH 2/5] update --- cosyvoice/bin/export_trt.py | 26 +++++++++++++++--------- cosyvoice/cli/cosyvoice.py | 5 ++++- cosyvoice/flow/flow.py | 2 +- cosyvoice/flow/flow_matching.py | 36 +++++++++++++++++++++++++++++++++ 4 files changed, 58 insertions(+), 11 deletions(-) diff --git a/cosyvoice/bin/export_trt.py b/cosyvoice/bin/export_trt.py index fea8205..0de874a 100644 --- a/cosyvoice/bin/export_trt.py +++ b/cosyvoice/bin/export_trt.py @@ -38,23 +38,21 @@ def main(): args = get_args() cosyvoice = CosyVoice(args.model_dir, load_jit=False, load_trt=False) - - flow = cosyvoice.model.flow estimator = cosyvoice.model.flow.decoder.estimator dtype = torch.float32 if not args.export_half else torch.float16 device = torch.device("cuda") batch_size = 1 - seq_len = 1024 - hidden_size = flow.output_size + seq_len = 256 + hidden_size = cosyvoice.model.flow.output_size x = torch.rand((batch_size, hidden_size, seq_len), dtype=dtype, device=device) - mask = torch.zeros((batch_size, 1, seq_len), dtype=dtype, device=device) + mask = torch.ones((batch_size, 1, seq_len), dtype=dtype, device=device) mu = torch.rand((batch_size, hidden_size, seq_len), dtype=dtype, device=device) - t = torch.tensor([0.], dtype=dtype, device=device) + t = torch.rand((batch_size, ), dtype=dtype, device=device) spks = torch.rand((batch_size, hidden_size), dtype=dtype, device=device) cond = torch.rand((batch_size, hidden_size, seq_len), dtype=dtype, device=device) - onnx_file_name = 'estimator_fp16.onnx' if args.export_half else 'estimator_fp32.onnx' + onnx_file_name = 'estimator_fp32.onnx' if not args.export_half else 'estimator_fp16.onnx' onnx_file_path = os.path.join(args.model_dir, onnx_file_name) dummy_input = (x, mask, mu, t, spks, cond) @@ -90,14 +88,24 @@ def main(): print(f"Adding TensorRT lib path {trt_lib_path} to LD_LIBRARY_PATH.") os.environ['LD_LIBRARY_PATH'] = f"{os.environ.get('LD_LIBRARY_PATH', '')}:{trt_lib_path}" - trt_file_name = 'estimator_fp16.plan' if args.export_half else 'estimator_fp32.plan' + trt_file_name = 'estimator_fp32.plan' if not args.export_half else 'estimator_fp16.plan' trt_file_path = os.path.join(args.model_dir, trt_file_name) trtexec_cmd = f"{tensorrt_path}/bin/trtexec --onnx={onnx_file_path} --saveEngine={trt_file_path} " \ "--minShapes=x:1x80x1,mask:1x1x1,mu:1x80x1,t:1,spks:1x80,cond:1x80x1 " \ - "--maxShapes=x:1x80x4096,mask:1x1x4096,mu:1x80x4096,t:1,spks:1x80,cond:1x80x4096 --verbose" + "--maxShapes=x:1x80x4096,mask:1x1x4096,mu:1x80x4096,t:1,spks:1x80,cond:1x80x4096 --verbose " + \ + ("--fp16" if args.export_half else "") +# /ossfs/workspace/TensorRT-10.2.0.19/bin/trtexec --onnx=estimator_fp32.onnx --saveEngine=estimator_fp32.plan --minShapes=x:1x80x1,mask:1x1x1,mu:1x80x1,t:1,spks:1x80,cond:1x80x1 --maxShapes=x:1x80x4096,mask:1x1x4096,mu:1x80x4096,t:1,spks:1x80,cond:1x80x4096 --verbose + print("execute ", trtexec_cmd) os.system(trtexec_cmd) + print("x.shape", x.shape) + print("mask.shape", mask.shape) + print("mu.shape", mu.shape) + print("t.shape", t.shape) + print("spks.shape", spks.shape) + print("cond.shape", cond.shape) + if __name__ == "__main__": main() diff --git a/cosyvoice/cli/cosyvoice.py b/cosyvoice/cli/cosyvoice.py index eac8922..3558374 100644 --- a/cosyvoice/cli/cosyvoice.py +++ b/cosyvoice/cli/cosyvoice.py @@ -21,7 +21,7 @@ from cosyvoice.utils.file_utils import logging class CosyVoice: - def __init__(self, model_dir, load_jit=True, load_trt=True, use_fp16=False): + def __init__(self, model_dir, load_jit=True, load_trt=False, use_fp16=False): instruct = True if '-Instruct' in model_dir else False self.model_dir = model_dir if not os.path.exists(model_dir): @@ -39,11 +39,14 @@ class CosyVoice: self.model.load('{}/llm.pt'.format(model_dir), '{}/flow.pt'.format(model_dir), '{}/hift.pt'.format(model_dir)) + load_jit = False if load_jit: self.model.load_jit('{}/llm.text_encoder.fp16.zip'.format(model_dir), '{}/llm.llm.fp16.zip'.format(model_dir)) + if load_trt: self.model.load_trt(model_dir, use_fp16) + del configs def list_avaliable_spks(self): diff --git a/cosyvoice/flow/flow.py b/cosyvoice/flow/flow.py index 5466542..4a40eed 100644 --- a/cosyvoice/flow/flow.py +++ b/cosyvoice/flow/flow.py @@ -107,7 +107,7 @@ class MaskedDiffWithXvec(torch.nn.Module): # concat text and prompt_text token_len1, token_len2 = prompt_token.shape[1], token.shape[1] token, token_len = torch.concat([prompt_token, token], dim=1), prompt_token_len + token_len - mask = (~make_pad_mask(token_len)).float().unsqueeze(-1).to(embedding) + mask = (~make_pad_mask(token_len)).to(embedding.dtype).unsqueeze(-1).to(embedding) token = self.input_embedding(torch.clamp(token, min=0)) * mask # text encode diff --git a/cosyvoice/flow/flow_matching.py b/cosyvoice/flow/flow_matching.py index 8cab545..ef024dc 100755 --- a/cosyvoice/flow/flow_matching.py +++ b/cosyvoice/flow/flow_matching.py @@ -32,6 +32,7 @@ class ConditionalCFM(BASECFM): self.estimator_context = None self.estimator_engine = None + self.is_saved = None @torch.inference_mode() def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None): @@ -123,6 +124,41 @@ class ConditionalCFM(BASECFM): self.estimator_context.execute_async_v3(stream_handle=handle) return ret else: + + if self.is_saved == None: + self.is_saved = True + output = self.estimator.forward(x, mask, mu, t, spks, cond) + torch.save(x, "x.pt") + torch.save(mask, "mask.pt") + torch.save(mu, "mu.pt") + torch.save(t, "t.pt") + torch.save(spks, "spks.pt") + torch.save(cond, "cond.pt") + torch.save(output, "output.pt") + dummy_input = (x, mask, mu, t, spks, cond) + torch.onnx.export( + self.estimator, + dummy_input, + "estimator_fp32.onnx", + export_params=True, + opset_version=17, + do_constant_folding=True, + input_names=['x', 'mask', 'mu', 't', 'spks', 'cond'], + output_names=['output'], + dynamic_axes={ + 'x': {2: 'seq_len'}, + 'mask': {2: 'seq_len'}, + 'mu': {2: 'seq_len'}, + 'cond': {2: 'seq_len'}, + 'output': {2: 'seq_len'}, + } + ) + # print("x, x.shape", x, x.shape) + # print("mask, mask.shape", mask, mask.shape) + # print("mu, mu.shape", mu, mu.shape) + # print("t, t.shape", t, t.shape) + # print("spks, spks.shape", spks, spks.shape) + # print("cond, cond.shape", cond, cond.shape) return self.estimator.forward(x, mask, mu, t, spks, cond) def compute_loss(self, x1, mask, mu, spks=None, cond=None): From 6e7f5b922a7ddca2d9537629fa624b997443970d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=BE=E6=81=AF?= Date: Fri, 30 Aug 2024 13:14:44 +0800 Subject: [PATCH 3/5] update --- cosyvoice/bin/export_trt.py | 6 ++-- cosyvoice/cli/model.py | 3 +- cosyvoice/flow/flow_matching.py | 52 +++++---------------------------- 3 files changed, 12 insertions(+), 49 deletions(-) diff --git a/cosyvoice/bin/export_trt.py b/cosyvoice/bin/export_trt.py index 0de874a..769bdca 100644 --- a/cosyvoice/bin/export_trt.py +++ b/cosyvoice/bin/export_trt.py @@ -66,13 +66,13 @@ def main(): opset_version=18, do_constant_folding=True, input_names=['x', 'mask', 'mu', 't', 'spks', 'cond'], - output_names=['output'], + output_names=['estimator_out'], dynamic_axes={ 'x': {2: 'seq_len'}, 'mask': {2: 'seq_len'}, 'mu': {2: 'seq_len'}, 'cond': {2: 'seq_len'}, - 'output': {2: 'seq_len'}, + 'estimator_out': {2: 'seq_len'}, } ) @@ -95,7 +95,7 @@ def main(): "--minShapes=x:1x80x1,mask:1x1x1,mu:1x80x1,t:1,spks:1x80,cond:1x80x1 " \ "--maxShapes=x:1x80x4096,mask:1x1x4096,mu:1x80x4096,t:1,spks:1x80,cond:1x80x4096 --verbose " + \ ("--fp16" if args.export_half else "") -# /ossfs/workspace/TensorRT-10.2.0.19/bin/trtexec --onnx=estimator_fp32.onnx --saveEngine=estimator_fp32.plan --minShapes=x:1x80x1,mask:1x1x1,mu:1x80x1,t:1,spks:1x80,cond:1x80x1 --maxShapes=x:1x80x4096,mask:1x1x4096,mu:1x80x4096,t:1,spks:1x80,cond:1x80x4096 --verbose + print("execute ", trtexec_cmd) os.system(trtexec_cmd) diff --git a/cosyvoice/cli/model.py b/cosyvoice/cli/model.py index f074fd8..59006a7 100644 --- a/cosyvoice/cli/model.py +++ b/cosyvoice/cli/model.py @@ -83,8 +83,7 @@ class CosyVoiceModel: with open(trt_file_path, 'rb') as f: serialized_engine = f.read() engine = runtime.deserialize_cuda_engine(serialized_engine) - self.flow.decoder.estimator_context = engine.create_execution_context() - self.flow.decoder.estimator_engine = engine + self.flow.decoder.estimator = engine.create_execution_context() def llm_job(self, text, prompt_text, llm_prompt_speech_token, llm_embedding, uuid): with self.llm_context: diff --git a/cosyvoice/flow/flow_matching.py b/cosyvoice/flow/flow_matching.py index ef024dc..8908d5f 100755 --- a/cosyvoice/flow/flow_matching.py +++ b/cosyvoice/flow/flow_matching.py @@ -30,10 +30,6 @@ class ConditionalCFM(BASECFM): # Just change the architecture of the estimator here self.estimator = estimator - self.estimator_context = None - self.estimator_engine = None - self.is_saved = None - @torch.inference_mode() def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None): """Forward diffusion @@ -102,7 +98,11 @@ class ConditionalCFM(BASECFM): return sol[-1] def forward_estimator(self, x, mask, mu, t, spks, cond): - if self.estimator_context is not None: + + if not isinstance(self.estimator, torch.nn.Module): + return self.estimator.forward(x, mask, mu, t, spks, cond) + + else: assert self.training is False, 'tensorrt cannot be used in training' bs = x.shape[0] hs = x.shape[1] @@ -116,50 +116,14 @@ class ConditionalCFM(BASECFM): self.estimator_context.set_input_shape("spks", spks.shape) self.estimator_context.set_input_shape("cond", cond.shape) bindings = [x.data_ptr(), mask.data_ptr(), mu.data_ptr(), t.data_ptr(), spks.data_ptr(), cond.data_ptr(), ret.data_ptr()] + names = ['x', 'mask', 'mu', 't', 'spks', 'cond', 'estimator_out'] for i in range(len(bindings)): - self.estimator_context.set_tensor_address(self.estimator_engine.get_tensor_name(i), bindings[i]) + self.estimator.set_tensor_address(names[i], bindings[i]) handle = torch.cuda.current_stream().cuda_stream - self.estimator_context.execute_async_v3(stream_handle=handle) + self.estimator.execute_async_v3(stream_handle=handle) return ret - else: - - if self.is_saved == None: - self.is_saved = True - output = self.estimator.forward(x, mask, mu, t, spks, cond) - torch.save(x, "x.pt") - torch.save(mask, "mask.pt") - torch.save(mu, "mu.pt") - torch.save(t, "t.pt") - torch.save(spks, "spks.pt") - torch.save(cond, "cond.pt") - torch.save(output, "output.pt") - dummy_input = (x, mask, mu, t, spks, cond) - torch.onnx.export( - self.estimator, - dummy_input, - "estimator_fp32.onnx", - export_params=True, - opset_version=17, - do_constant_folding=True, - input_names=['x', 'mask', 'mu', 't', 'spks', 'cond'], - output_names=['output'], - dynamic_axes={ - 'x': {2: 'seq_len'}, - 'mask': {2: 'seq_len'}, - 'mu': {2: 'seq_len'}, - 'cond': {2: 'seq_len'}, - 'output': {2: 'seq_len'}, - } - ) - # print("x, x.shape", x, x.shape) - # print("mask, mask.shape", mask, mask.shape) - # print("mu, mu.shape", mu, mu.shape) - # print("t, t.shape", t, t.shape) - # print("spks, spks.shape", spks, spks.shape) - # print("cond, cond.shape", cond, cond.shape) - return self.estimator.forward(x, mask, mu, t, spks, cond) def compute_loss(self, x1, mask, mu, spks=None, cond=None): """Computes diffusion loss From 29408360fb4484b6ae53810b181ceec9b79611d9 Mon Sep 17 00:00:00 2001 From: "zhoubofan.zbf" Date: Fri, 30 Aug 2024 13:43:54 +0800 Subject: [PATCH 4/5] fix bug --- cosyvoice/bin/export_trt.py | 19 ++++++++++--------- cosyvoice/cli/cosyvoice.py | 4 ++-- cosyvoice/cli/model.py | 3 ++- cosyvoice/flow/flow_matching.py | 8 ++++---- 4 files changed, 18 insertions(+), 16 deletions(-) diff --git a/cosyvoice/bin/export_trt.py b/cosyvoice/bin/export_trt.py index 769bdca..1bf958c 100644 --- a/cosyvoice/bin/export_trt.py +++ b/cosyvoice/bin/export_trt.py @@ -11,7 +11,7 @@ except ImportError: error_msg_zh = [ "step.1 下载 tensorrt .tar.gz 压缩包并解压，下载地址: https://developer.nvidia.com/tensorrt/download/10x", "step.2 使用 tensorrt whl 包进行安装根据 python 版本对应进行安装，如 pip install ${TensorRT-Path}/python/tensorrt-10.2.0-cp38-none-linux_x86_64.whl", - "step.3 将 tensorrt 的 lib 路径添加进环境变量中，export LD_LIBRARY_PATH=${TensorRT-Path}/lib/" + "step.3 将 tensorrt 的 lib 路径添加进环境变量中，export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${TensorRT-Path}/lib/" ] print("\n".join(error_msg_zh)) sys.exit(1) @@ -23,7 +23,7 @@ def get_args(): parser = argparse.ArgumentParser(description='Export your model for deployment') parser.add_argument('--model_dir', type=str, - default='pretrained_models/CosyVoice-300M', + default='pretrained_models/CosyVoice-300M-SFT', help='Local path to the model directory') parser.add_argument('--export_half', @@ -91,7 +91,8 @@ def main(): trt_file_name = 'estimator_fp32.plan' if not args.export_half else 'estimator_fp16.plan' trt_file_path = os.path.join(args.model_dir, trt_file_name) - trtexec_cmd = f"{tensorrt_path}/bin/trtexec --onnx={onnx_file_path} --saveEngine={trt_file_path} " \ + trtexec_bin = os.path.join(tensorrt_path, 'bin/trtexec') + trtexec_cmd = f"{trtexec_bin} --onnx={onnx_file_path} --saveEngine={trt_file_path} " \ "--minShapes=x:1x80x1,mask:1x1x1,mu:1x80x1,t:1,spks:1x80,cond:1x80x1 " \ "--maxShapes=x:1x80x4096,mask:1x1x4096,mu:1x80x4096,t:1,spks:1x80,cond:1x80x4096 --verbose " + \ ("--fp16" if args.export_half else "") @@ -100,12 +101,12 @@ def main(): os.system(trtexec_cmd) - print("x.shape", x.shape) - print("mask.shape", mask.shape) - print("mu.shape", mu.shape) - print("t.shape", t.shape) - print("spks.shape", spks.shape) - print("cond.shape", cond.shape) + # print("x.shape", x.shape) + # print("mask.shape", mask.shape) + # print("mu.shape", mu.shape) + # print("t.shape", t.shape) + # print("spks.shape", spks.shape) + # print("cond.shape", cond.shape) if __name__ == "__main__": main() diff --git a/cosyvoice/cli/cosyvoice.py b/cosyvoice/cli/cosyvoice.py index 3558374..5028ad1 100644 --- a/cosyvoice/cli/cosyvoice.py +++ b/cosyvoice/cli/cosyvoice.py @@ -21,7 +21,7 @@ from cosyvoice.utils.file_utils import logging class CosyVoice: - def __init__(self, model_dir, load_jit=True, load_trt=False, use_fp16=False): + def __init__(self, model_dir, load_jit=True, load_trt=True, use_fp16=False): instruct = True if '-Instruct' in model_dir else False self.model_dir = model_dir if not os.path.exists(model_dir): @@ -39,7 +39,7 @@ class CosyVoice: self.model.load('{}/llm.pt'.format(model_dir), '{}/flow.pt'.format(model_dir), '{}/hift.pt'.format(model_dir)) - load_jit = False + if load_jit: self.model.load_jit('{}/llm.text_encoder.fp16.zip'.format(model_dir), '{}/llm.llm.fp16.zip'.format(model_dir)) diff --git a/cosyvoice/cli/model.py b/cosyvoice/cli/model.py index 59006a7..bf18ea3 100644 --- a/cosyvoice/cli/model.py +++ b/cosyvoice/cli/model.py @@ -83,7 +83,8 @@ class CosyVoiceModel: with open(trt_file_path, 'rb') as f: serialized_engine = f.read() engine = runtime.deserialize_cuda_engine(serialized_engine) - self.flow.decoder.estimator = engine.create_execution_context() + self.flow.decoder.estimator_context = engine.create_execution_context() + self.flow.decoder.estimator = None def llm_job(self, text, prompt_text, llm_prompt_speech_token, llm_embedding, uuid): with self.llm_context: diff --git a/cosyvoice/flow/flow_matching.py b/cosyvoice/flow/flow_matching.py index 8908d5f..18efe75 100755 --- a/cosyvoice/flow/flow_matching.py +++ b/cosyvoice/flow/flow_matching.py @@ -99,10 +99,10 @@ class ConditionalCFM(BASECFM): def forward_estimator(self, x, mask, mu, t, spks, cond): - if not isinstance(self.estimator, torch.nn.Module): + if self.estimator is not None: return self.estimator.forward(x, mask, mu, t, spks, cond) - else: + print("-----------") assert self.training is False, 'tensorrt cannot be used in training' bs = x.shape[0] hs = x.shape[1] @@ -119,10 +119,10 @@ class ConditionalCFM(BASECFM): names = ['x', 'mask', 'mu', 't', 'spks', 'cond', 'estimator_out'] for i in range(len(bindings)): - self.estimator.set_tensor_address(names[i], bindings[i]) + self.estimator_context.set_tensor_address(names[i], bindings[i]) handle = torch.cuda.current_stream().cuda_stream - self.estimator.execute_async_v3(stream_handle=handle) + self.estimator_context.execute_async_v3(stream_handle=handle) return ret def compute_loss(self, x1, mask, mu, spks=None, cond=None): From 18599be8d5edba918fcaac9e0f55ecd6ab4e045c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=BE=E6=81=AF?= Date: Fri, 30 Aug 2024 14:15:24 +0800 Subject: [PATCH 5/5] mirror modify --- cosyvoice/bin/export_trt.py | 14 ++++++++++++++ cosyvoice/cli/model.py | 6 +----- cosyvoice/flow/flow_matching.py | 1 - 3 files changed, 15 insertions(+), 6 deletions(-) diff --git a/cosyvoice/bin/export_trt.py b/cosyvoice/bin/export_trt.py index 1bf958c..c737373 100644 --- a/cosyvoice/bin/export_trt.py +++ b/cosyvoice/bin/export_trt.py @@ -1,3 +1,17 @@ +# Copyright (c) 2024 Antgroup Inc (authors: Zhoubofan, hexisyztem@icloud.com) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import argparse import logging import os diff --git a/cosyvoice/cli/model.py b/cosyvoice/cli/model.py index bf18ea3..50ae0b1 100644 --- a/cosyvoice/cli/model.py +++ b/cosyvoice/cli/model.py @@ -20,11 +20,6 @@ from contextlib import nullcontext import uuid from cosyvoice.utils.common import fade_in_out -try: - import tensorrt as trt -except ImportError: - ... - class CosyVoiceModel: def __init__(self, @@ -72,6 +67,7 @@ class CosyVoiceModel: self.llm.llm = llm_llm def load_trt(self, model_dir, use_fp16): + import tensorrt as trt trt_file_name = 'estimator_fp16.plan' if use_fp16 else 'estimator_fp32.plan' trt_file_path = os.path.join(model_dir, trt_file_name) if not os.path.isfile(trt_file_path): diff --git a/cosyvoice/flow/flow_matching.py b/cosyvoice/flow/flow_matching.py index 18efe75..a31506a 100755 --- a/cosyvoice/flow/flow_matching.py +++ b/cosyvoice/flow/flow_matching.py @@ -102,7 +102,6 @@ class ConditionalCFM(BASECFM): if self.estimator is not None: return self.estimator.forward(x, mask, mu, t, spks, cond) else: - print("-----------") assert self.training is False, 'tensorrt cannot be used in training' bs = x.shape[0] hs = x.shape[1]