mirror modify

cosyvoice/bin/export_trt.py

@@ -1,3 +1,17 @@
+# Copyright (c) 2024 Antgroup Inc (authors: Zhoubofan, [email protected])
+#
+# 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

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

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]