mirror of
https://github.com/FunAudioLLM/CosyVoice.git
synced 2026-02-05 09:59:23 +08:00
add llm export script
This commit is contained in:
@@ -21,7 +21,7 @@ from cosyvoice.utils.file_utils import logging
|
||||
|
||||
class CosyVoice:
|
||||
|
||||
def __init__(self, model_dir):
|
||||
def __init__(self, model_dir, load_script=True):
|
||||
instruct = True if '-Instruct' in model_dir else False
|
||||
self.model_dir = model_dir
|
||||
if not os.path.exists(model_dir):
|
||||
@@ -39,6 +39,9 @@ class CosyVoice:
|
||||
self.model.load('{}/llm.pt'.format(model_dir),
|
||||
'{}/flow.pt'.format(model_dir),
|
||||
'{}/hift.pt'.format(model_dir))
|
||||
if load_script:
|
||||
self.model.load_script('{}/llm.text_encoder.fp16.zip'.format(model_dir),
|
||||
'{}/llm.llm.fp16.zip'.format(model_dir))
|
||||
del configs
|
||||
|
||||
def list_avaliable_spks(self):
|
||||
|
||||
@@ -47,11 +47,18 @@ class CosyVoiceModel:
|
||||
def load(self, llm_model, flow_model, hift_model):
|
||||
self.llm.load_state_dict(torch.load(llm_model, map_location=self.device))
|
||||
self.llm.to(self.device).eval()
|
||||
self.llm.half()
|
||||
self.flow.load_state_dict(torch.load(flow_model, map_location=self.device))
|
||||
self.flow.to(self.device).eval()
|
||||
self.hift.load_state_dict(torch.load(hift_model, map_location=self.device))
|
||||
self.hift.to(self.device).eval()
|
||||
|
||||
def load_script(self, llm_text_encoder_model, llm_llm_model):
|
||||
llm_text_encoder = torch.jit.load(llm_text_encoder_model)
|
||||
self.llm.text_encoder = llm_text_encoder
|
||||
llm_llm = torch.jit.load(llm_llm_model)
|
||||
self.llm.llm = llm_llm
|
||||
|
||||
def llm_job(self, text, text_len, prompt_text, prompt_text_len, llm_prompt_speech_token, llm_prompt_speech_token_len, llm_embedding, this_uuid):
|
||||
with self.llm_context:
|
||||
for i in self.llm.inference(text=text.to(self.device),
|
||||
@@ -60,7 +67,7 @@ class CosyVoiceModel:
|
||||
prompt_text_len=prompt_text_len.to(self.device),
|
||||
prompt_speech_token=llm_prompt_speech_token.to(self.device),
|
||||
prompt_speech_token_len=llm_prompt_speech_token_len.to(self.device),
|
||||
embedding=llm_embedding.to(self.device),
|
||||
embedding=llm_embedding.to(self.device).half(),
|
||||
sampling=25,
|
||||
max_token_text_ratio=30,
|
||||
min_token_text_ratio=3):
|
||||
|
||||
Reference in New Issue
Block a user