add func export_codec_vllm

This commit is contained in:
雾聪
2025-02-25 15:01:29 +08:00
parent fd45708e4b
commit c37c00ff94
2 changed files with 36 additions and 1 deletions

View File

@@ -126,7 +126,7 @@ class CosyVoice:
class CosyVoice2(CosyVoice):
def __init__(self, model_dir, load_jit=False, load_trt=False, fp16=False):
def __init__(self, model_dir, load_jit=False, load_trt=False, fp16=False, use_vllm=False):
self.instruct = True if '-Instruct' in model_dir else False
self.model_dir = model_dir
self.fp16 = fp16
@@ -149,6 +149,8 @@ class CosyVoice2(CosyVoice):
self.model.load('{}/llm.pt'.format(model_dir),
'{}/flow.pt'.format(model_dir),
'{}/hift.pt'.format(model_dir))
if use_vllm:
self.model.export_codec_vllm(''.join([model_dir, '/codec_vllm_model']))
if load_jit:
self.model.load_jit('{}/flow.encoder.{}.zip'.format(model_dir, 'fp16' if self.fp16 is True else 'fp32'))
if load_trt: