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):