deepgeek/CosyVoice
2
0
Fork 0
mirror of https://github.com/FunAudioLLM/CosyVoice.git synced 2026-02-04 09:29:25 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix bug

Browse Source
This commit is contained in:
zhoubofan.zbf
2024-08-30 13:43:54 +08:00
parent 6e7f5b922a
commit 29408360fb
4 changed files with 18 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
cosyvoice/bin/export_trt.py
View File

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

4
cosyvoice/cli/cosyvoice.py
View File

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

3
cosyvoice/cli/model.py
View File

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

8
cosyvoice/flow/flow_matching.py
View File

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