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

add onnx export

Browse Source
This commit is contained in:
lyuxiang.lx
2024-09-04 18:15:33 +08:00
parent d8197de4cc
commit 2ce724045b
6 changed files with 105 additions and 280 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
cosyvoice/bin/export_jit.py
View File

@@ -44,7 +44,7 @@ def main():
torch._C._jit_set_profiling_mode(False)
torch._C._jit_set_profiling_executor(False)
cosyvoice = CosyVoice(args.model_dir, load_jit=False, load_trt=False)
cosyvoice = CosyVoice(args.model_dir, load_jit=False, load_onnx=False)
# 1. export llm text_encoder
llm_text_encoder = cosyvoice.model.llm.text_encoder.half()
@@ -60,5 +60,12 @@ def main():
script = torch.jit.optimize_for_inference(script)
script.save('{}/llm.llm.fp16.zip'.format(args.model_dir))
# 3. export flow encoder
flow_encoder = cosyvoice.model.flow.encoder
script = torch.jit.script(flow_encoder)
script = torch.jit.freeze(script)
script = torch.jit.optimize_for_inference(script)
script.save('{}/flow.encoder.fp32.zip'.format(args.model_dir))
if __name__ == '__main__':
main()