fix bug

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

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

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:

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