From 62e04e8856402c90a3af29b24b348c31ab506f46 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=BE=E6=81=AF?= Date: Wed, 16 Apr 2025 11:16:28 +0800 Subject: [PATCH 1/4] Enhance CosyVoice with CUDA stream management and estimator handling - Introduced a queue-based system for managing CUDA streams to improve inference performance. - Updated inference methods to utilize CUDA streams for asynchronous processing. - Added an EstimatorWrapper class to manage TensorRT estimators, allowing for efficient execution context handling. - Modified model loading functions to support estimator count configuration. - Improved logging and performance tracking during inference operations. --- cosyvoice/cli/cosyvoice.py | 200 ++++++++++++++++++++------------ cosyvoice/cli/model.py | 43 +++---- cosyvoice/flow/flow_matching.py | 77 +++++++++--- 3 files changed, 207 insertions(+), 113 deletions(-) diff --git a/cosyvoice/cli/cosyvoice.py b/cosyvoice/cli/cosyvoice.py index 39464ca..f49d2da 100644 --- a/cosyvoice/cli/cosyvoice.py +++ b/cosyvoice/cli/cosyvoice.py @@ -22,7 +22,7 @@ from cosyvoice.cli.frontend import CosyVoiceFrontEnd from cosyvoice.cli.model import CosyVoiceModel, CosyVoice2Model, VllmCosyVoice2Model from cosyvoice.utils.file_utils import logging from cosyvoice.utils.class_utils import get_model_type - +import queue class CosyVoice: @@ -54,11 +54,18 @@ class CosyVoice: '{}/llm.llm.{}.zip'.format(model_dir, 'fp16' if self.fp16 is True else 'fp32'), '{}/flow.encoder.{}.zip'.format(model_dir, 'fp16' if self.fp16 is True else 'fp32')) if load_trt: + self.estimator_count = configs['flow']['decoder']['estimator'].get('estimator_count', 1) self.model.load_trt('{}/flow.decoder.estimator.{}.mygpu.plan'.format(model_dir, 'fp16' if self.fp16 is True else 'fp32'), '{}/flow.decoder.estimator.fp32.onnx'.format(model_dir), - self.fp16) + self.fp16, self.estimator_count) del configs + thread_count = 10 + self.stream_pool = queue.Queue(maxsize=thread_count) + for _ in range(thread_count): + self.stream_pool.put(torch.cuda.Stream(self.device)) + + def list_available_spks(self): spks = list(self.frontend.spk2info.keys()) return spks @@ -67,80 +74,104 @@ class CosyVoice: self.frontend.add_spk_info(spk_id, spk_info) def inference_sft(self, tts_text, spk_id, stream=False, speed=1.0, text_frontend=True): - for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): - model_input = self.frontend.frontend_sft(i, spk_id) - start_time = time.time() - logging.info('synthesis text {}'.format(i)) - for model_output in self.model.tts(**model_input, stream=stream, speed=speed): - speech_len = model_output['tts_speech'].shape[1] / self.sample_rate - logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) - yield model_output + cuda_stream = self.stream_pool.get() + with torch.cuda.stream(cuda_stream): + for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): + model_input = self.frontend.frontend_sft(i, spk_id) start_time = time.time() + logging.info('synthesis text {}'.format(i)) + for model_output in self.model.tts(**model_input, stream=stream, speed=speed): + speech_len = model_output['tts_speech'].shape[1] / self.sample_rate + logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) + yield model_output + start_time = time.time() + cuda_stream.synchronize() + self.stream_pool.put(cuda_stream) def inference_zero_shot(self, tts_text, prompt_text, prompt_speech_16k, stream=False, speed=1.0, text_frontend=True): - prompt_text = self.frontend.text_normalize(prompt_text, split=False, text_frontend=text_frontend) - for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): - if (not isinstance(i, Generator)) and len(i) < 0.5 * len(prompt_text): - logging.warning('synthesis text {} too short than prompt text {}, this may lead to bad performance'.format(i, prompt_text)) - model_input = self.frontend.frontend_zero_shot(i, prompt_text, prompt_speech_16k, self.sample_rate) - start_time = time.time() - logging.info('synthesis text {}'.format(i)) - for model_output in self.model.tts(**model_input, stream=stream, speed=speed): - speech_len = model_output['tts_speech'].shape[1] / self.sample_rate - logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) - yield model_output + cuda_stream = self.stream_pool.get() + with torch.cuda.stream(cuda_stream): + prompt_text = self.frontend.text_normalize(prompt_text, split=False, text_frontend=text_frontend) + for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): + if (not isinstance(i, Generator)) and len(i) < 0.5 * len(prompt_text): + logging.warning('synthesis text {} too short than prompt text {}, this may lead to bad performance'.format(i, prompt_text)) + model_input = self.frontend.frontend_zero_shot(i, prompt_text, prompt_speech_16k, self.sample_rate) start_time = time.time() + logging.info('synthesis text {}'.format(i)) + for model_output in self.model.tts(**model_input, stream=stream, speed=speed): + speech_len = model_output['tts_speech'].shape[1] / self.sample_rate + logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) + yield model_output + start_time = time.time() + cuda_stream.synchronize() + self.stream_pool.put(cuda_stream) def inference_zero_shot_by_spk_id(self, tts_text, spk_id, stream=False, speed=1.0, text_frontend=True): """使用预定义的说话人执行 zero_shot 推理""" - for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): - model_input = self.frontend.frontend_zero_shot_by_spk_id(i, spk_id) - start_time = time.time() - last_time = start_time - chunk_index = 0 - logging.info('synthesis text {}'.format(i)) - for model_output in self.model.tts(**model_input, stream=stream, speed=speed): - speech_len = model_output['tts_speech'].shape[1] / self.sample_rate - logging.info('yield speech index:{}, len {:.2f}, rtf {:.3f}, cost {:.3f}s, all cost time {:.3f}s'.format( - chunk_index, speech_len, (time.time()-last_time)/speech_len, time.time()-last_time, time.time()-start_time)) - yield model_output - last_time = time.time() - chunk_index += 1 + cuda_stream = self.stream_pool.get() + with torch.cuda.stream(cuda_stream): + for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): + model_input = self.frontend.frontend_zero_shot_by_spk_id(i, spk_id) + start_time = time.time() + last_time = start_time + chunk_index = 0 + logging.info('synthesis text {}'.format(i)) + for model_output in self.model.tts(**model_input, stream=stream, speed=speed): + speech_len = model_output['tts_speech'].shape[1] / self.sample_rate + logging.info('yield speech index:{}, len {:.2f}, rtf {:.3f}, cost {:.3f}s, all cost time {:.3f}s'.format( + chunk_index, speech_len, (time.time()-last_time)/speech_len, time.time()-last_time, time.time()-start_time)) + yield model_output + last_time = time.time() + chunk_index += 1 + cuda_stream.synchronize() + self.stream_pool.put(cuda_stream) def inference_cross_lingual(self, tts_text, prompt_speech_16k, stream=False, speed=1.0, text_frontend=True): - for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): - model_input = self.frontend.frontend_cross_lingual(i, prompt_speech_16k, self.sample_rate) - start_time = time.time() - logging.info('synthesis text {}'.format(i)) - for model_output in self.model.tts(**model_input, stream=stream, speed=speed): - speech_len = model_output['tts_speech'].shape[1] / self.sample_rate - logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) - yield model_output + cuda_stream = self.stream_pool.get() + with torch.cuda.stream(cuda_stream): + for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): + model_input = self.frontend.frontend_cross_lingual(i, prompt_speech_16k, self.sample_rate) start_time = time.time() + logging.info('synthesis text {}'.format(i)) + for model_output in self.model.tts(**model_input, stream=stream, speed=speed): + speech_len = model_output['tts_speech'].shape[1] / self.sample_rate + logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) + yield model_output + start_time = time.time() + cuda_stream.synchronize() + self.stream_pool.put(cuda_stream) def inference_instruct(self, tts_text, spk_id, instruct_text, stream=False, speed=1.0, text_frontend=True): - assert isinstance(self.model, CosyVoiceModel), 'inference_instruct is only implemented for CosyVoice!' - if self.instruct is False: - raise ValueError('{} do not support instruct inference'.format(self.model_dir)) - instruct_text = self.frontend.text_normalize(instruct_text, split=False, text_frontend=text_frontend) - for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): - model_input = self.frontend.frontend_instruct(i, spk_id, instruct_text) + cuda_stream = self.stream_pool.get() + with torch.cuda.stream(cuda_stream): + assert isinstance(self.model, CosyVoiceModel), 'inference_instruct is only implemented for CosyVoice!' + if self.instruct is False: + raise ValueError('{} do not support instruct inference'.format(self.model_dir)) + instruct_text = self.frontend.text_normalize(instruct_text, split=False, text_frontend=text_frontend) + for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): + model_input = self.frontend.frontend_instruct(i, spk_id, instruct_text) + start_time = time.time() + logging.info('synthesis text {}'.format(i)) + for model_output in self.model.tts(**model_input, stream=stream, speed=speed): + speech_len = model_output['tts_speech'].shape[1] / self.sample_rate + logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) + yield model_output + start_time = time.time() + cuda_stream.synchronize() + self.stream_pool.put(cuda_stream) + + def inference_vc(self, source_speech_16k, prompt_speech_16k, stream=False, speed=1.0): + cuda_stream = self.stream_pool.get() + with torch.cuda.stream(cuda_stream): + model_input = self.frontend.frontend_vc(source_speech_16k, prompt_speech_16k, self.sample_rate) start_time = time.time() - logging.info('synthesis text {}'.format(i)) - for model_output in self.model.tts(**model_input, stream=stream, speed=speed): + for model_output in self.model.vc(**model_input, stream=stream, speed=speed): speech_len = model_output['tts_speech'].shape[1] / self.sample_rate logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) yield model_output start_time = time.time() - - def inference_vc(self, source_speech_16k, prompt_speech_16k, stream=False, speed=1.0): - model_input = self.frontend.frontend_vc(source_speech_16k, prompt_speech_16k, self.sample_rate) - start_time = time.time() - for model_output in self.model.vc(**model_input, stream=stream, speed=speed): - speech_len = model_output['tts_speech'].shape[1] / self.sample_rate - logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) - yield model_output - start_time = time.time() + cuda_stream.synchronize() + self.stream_pool.put(cuda_stream) class CosyVoice2(CosyVoice): @@ -178,33 +209,48 @@ class CosyVoice2(CosyVoice): if load_jit: self.model.load_jit('{}/flow.encoder.{}.zip'.format(model_dir, 'fp16' if self.fp16 is True else 'fp32')) if load_trt: + self.estimator_count = configs['flow']['decoder']['estimator'].get('estimator_count', 1) self.model.load_trt('{}/flow.decoder.estimator.{}.mygpu.plan'.format(model_dir, 'fp16' if self.fp16 is True else 'fp32'), '{}/flow.decoder.estimator.fp32.onnx'.format(model_dir), - self.fp16) + self.fp16, self.estimator_count) del configs + thread_count = 10 + self.stream_pool = queue.Queue(maxsize=thread_count) + for _ in range(thread_count): + self.stream_pool.put(torch.cuda.Stream(self.device)) + def inference_instruct(self, *args, **kwargs): raise NotImplementedError('inference_instruct is not implemented for CosyVoice2!') def inference_instruct2(self, tts_text, instruct_text, prompt_speech_16k, stream=False, speed=1.0, text_frontend=True): - assert isinstance(self.model, CosyVoice2Model), 'inference_instruct2 is only implemented for CosyVoice2!' - for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): - model_input = self.frontend.frontend_instruct2(i, instruct_text, prompt_speech_16k, self.sample_rate) - start_time = time.time() - logging.info('synthesis text {}'.format(i)) - for model_output in self.model.tts(**model_input, stream=stream, speed=speed): - speech_len = model_output['tts_speech'].shape[1] / self.sample_rate - logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) - yield model_output + cuda_stream = self.stream_pool.get() + with torch.cuda.stream(cuda_stream): + assert isinstance(self.model, CosyVoice2Model), 'inference_instruct2 is only implemented for CosyVoice2!' + for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): + model_input = self.frontend.frontend_instruct2(i, instruct_text, prompt_speech_16k, self.sample_rate) start_time = time.time() + logging.info('synthesis text {}'.format(i)) + for model_output in self.model.tts(**model_input, stream=stream, speed=speed): + speech_len = model_output['tts_speech'].shape[1] / self.sample_rate + logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) + yield model_output + start_time = time.time() + cuda_stream.synchronize() + self.stream_pool.put(cuda_stream) def inference_instruct2_by_spk_id(self, tts_text, instruct_text, spk_id, stream=False, speed=1.0, text_frontend=True): - for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): - model_input = self.frontend.frontend_instruct2_by_spk_id(i, instruct_text, spk_id) - start_time = time.time() - logging.info('synthesis text {}'.format(i)) - for model_output in self.model.tts(**model_input, stream=stream, speed=speed): - speech_len = model_output['tts_speech'].shape[1] / self.sample_rate - logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) - yield model_output + cuda_stream = self.stream_pool.get() + with torch.cuda.stream(cuda_stream): + assert isinstance(self.model, CosyVoice2Model), 'inference_instruct2 is only implemented for CosyVoice2!' + for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): + model_input = self.frontend.frontend_instruct2_by_spk_id(i, instruct_text, spk_id) start_time = time.time() + logging.info('synthesis text {}'.format(i)) + for model_output in self.model.tts(**model_input, stream=stream, speed=speed): + speech_len = model_output['tts_speech'].shape[1] / self.sample_rate + logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) + yield model_output + start_time = time.time() + cuda_stream.synchronize() + self.stream_pool.put(cuda_stream) diff --git a/cosyvoice/cli/model.py b/cosyvoice/cli/model.py index c0d25ba..769dc92 100644 --- a/cosyvoice/cli/model.py +++ b/cosyvoice/cli/model.py @@ -22,7 +22,7 @@ from contextlib import nullcontext import uuid from cosyvoice.utils.common import fade_in_out from cosyvoice.utils.file_utils import convert_onnx_to_trt - +from cosyvoice.flow.flow_matching import EstimatorWrapper class CosyVoiceModel: @@ -84,7 +84,7 @@ class CosyVoiceModel: flow_encoder = torch.jit.load(flow_encoder_model, map_location=self.device) self.flow.encoder = flow_encoder - def load_trt(self, flow_decoder_estimator_model, flow_decoder_onnx_model, fp16): + def load_trt(self, flow_decoder_estimator_model, flow_decoder_onnx_model, fp16, estimator_count=1): assert torch.cuda.is_available(), 'tensorrt only supports gpu!' if not os.path.exists(flow_decoder_estimator_model): convert_onnx_to_trt(flow_decoder_estimator_model, flow_decoder_onnx_model, fp16) @@ -96,7 +96,7 @@ class CosyVoiceModel: self.flow.decoder.estimator_engine = trt.Runtime(trt.Logger(trt.Logger.INFO)).deserialize_cuda_engine(f.read()) if self.flow.decoder.estimator_engine is None: raise ValueError('failed to load trt {}'.format(flow_decoder_estimator_model)) - self.flow.decoder.estimator = self.flow.decoder.estimator_engine.create_execution_context() + self.flow.decoder.estimator = EstimatorWrapper(self.flow.decoder.estimator_engine, estimator_count=estimator_count) def llm_job(self, text, prompt_text, llm_prompt_speech_token, llm_embedding, uuid): with self.llm_context: @@ -122,13 +122,13 @@ class CosyVoiceModel: def token2wav(self, token, prompt_token, prompt_feat, embedding, uuid, finalize=False, speed=1.0): tts_mel, flow_cache = self.flow.inference(token=token.to(self.device), - token_len=torch.tensor([token.shape[1]], dtype=torch.int32).to(self.device), - prompt_token=prompt_token.to(self.device), - prompt_token_len=torch.tensor([prompt_token.shape[1]], dtype=torch.int32).to(self.device), - prompt_feat=prompt_feat.to(self.device), - prompt_feat_len=torch.tensor([prompt_feat.shape[1]], dtype=torch.int32).to(self.device), - embedding=embedding.to(self.device), - flow_cache=self.flow_cache_dict[uuid]) + token_len=torch.tensor([token.shape[1]], dtype=torch.int32).to(self.device), + prompt_token=prompt_token.to(self.device), + prompt_token_len=torch.tensor([prompt_token.shape[1]], dtype=torch.int32).to(self.device), + prompt_feat=prompt_feat.to(self.device), + prompt_feat_len=torch.tensor([prompt_feat.shape[1]], dtype=torch.int32).to(self.device), + embedding=embedding.to(self.device), + flow_cache=self.flow_cache_dict[uuid]) self.flow_cache_dict[uuid] = flow_cache # mel overlap fade in out @@ -148,8 +148,8 @@ class CosyVoiceModel: if self.hift_cache_dict[uuid] is not None: tts_speech = fade_in_out(tts_speech, self.hift_cache_dict[uuid]['speech'], self.speech_window) self.hift_cache_dict[uuid] = {'mel': tts_mel[:, :, -self.mel_cache_len:], - 'source': tts_source[:, :, -self.source_cache_len:], - 'speech': tts_speech[:, -self.source_cache_len:]} + 'source': tts_source[:, :, -self.source_cache_len:], + 'speech': tts_speech[:, -self.source_cache_len:]} tts_speech = tts_speech[:, :-self.source_cache_len] else: if speed != 1.0: @@ -319,14 +319,15 @@ class CosyVoice2Model(CosyVoiceModel): self.flow.encoder = flow_encoder def token2wav(self, token, prompt_token, prompt_feat, embedding, uuid, token_offset, finalize=False, speed=1.0): + tts_mel, _ = self.flow.inference(token=token.to(self.device), - token_len=torch.tensor([token.shape[1]], dtype=torch.int32).to(self.device), - prompt_token=prompt_token.to(self.device), - prompt_token_len=torch.tensor([prompt_token.shape[1]], dtype=torch.int32).to(self.device), - prompt_feat=prompt_feat.to(self.device), - prompt_feat_len=torch.tensor([prompt_feat.shape[1]], dtype=torch.int32).to(self.device), - embedding=embedding.to(self.device), - finalize=finalize) + token_len=torch.tensor([token.shape[1]], dtype=torch.int32).to(self.device), + prompt_token=prompt_token.to(self.device), + prompt_token_len=torch.tensor([prompt_token.shape[1]], dtype=torch.int32).to(self.device), + prompt_feat=prompt_feat.to(self.device), + prompt_feat_len=torch.tensor([prompt_feat.shape[1]], dtype=torch.int32).to(self.device), + embedding=embedding.to(self.device), + finalize=finalize) tts_mel = tts_mel[:, :, token_offset * self.flow.token_mel_ratio:] # append hift cache if self.hift_cache_dict[uuid] is not None: @@ -340,8 +341,8 @@ class CosyVoice2Model(CosyVoiceModel): if self.hift_cache_dict[uuid] is not None: tts_speech = fade_in_out(tts_speech, self.hift_cache_dict[uuid]['speech'], self.speech_window) self.hift_cache_dict[uuid] = {'mel': tts_mel[:, :, -self.mel_cache_len:], - 'source': tts_source[:, :, -self.source_cache_len:], - 'speech': tts_speech[:, -self.source_cache_len:]} + 'source': tts_source[:, :, -self.source_cache_len:], + 'speech': tts_speech[:, -self.source_cache_len:]} tts_speech = tts_speech[:, :-self.source_cache_len] else: if speed != 1.0: diff --git a/cosyvoice/flow/flow_matching.py b/cosyvoice/flow/flow_matching.py index 6a60f6d..39643ed 100644 --- a/cosyvoice/flow/flow_matching.py +++ b/cosyvoice/flow/flow_matching.py @@ -15,7 +15,26 @@ import threading import torch import torch.nn.functional as F from matcha.models.components.flow_matching import BASECFM +import queue +class EstimatorWrapper: + def __init__(self, estimator_engine, estimator_count=2,): + self.estimators = queue.Queue() + self.estimator_engine = estimator_engine + for _ in range(estimator_count): + estimator = estimator_engine.create_execution_context() + if estimator is not None: + self.estimators.put(estimator) + + if self.estimators.empty(): + raise Exception("No available estimator") + + def acquire_estimator(self): + return self.estimators.get(), self.estimator_engine + + def release_estimator(self, estimator): + self.estimators.put(estimator) + return class ConditionalCFM(BASECFM): def __init__(self, in_channels, cfm_params, n_spks=1, spk_emb_dim=64, estimator: torch.nn.Module = None): @@ -125,22 +144,50 @@ class ConditionalCFM(BASECFM): if isinstance(self.estimator, torch.nn.Module): return self.estimator.forward(x, mask, mu, t, spks, cond) else: - with self.lock: - self.estimator.set_input_shape('x', (2, 80, x.size(2))) - self.estimator.set_input_shape('mask', (2, 1, x.size(2))) - self.estimator.set_input_shape('mu', (2, 80, x.size(2))) - self.estimator.set_input_shape('t', (2,)) - self.estimator.set_input_shape('spks', (2, 80)) - self.estimator.set_input_shape('cond', (2, 80, x.size(2))) + if isinstance(self.estimator, EstimatorWrapper): + estimator, engine = self.estimator.acquire_estimator() + + estimator.set_input_shape('x', (2, 80, x.size(2))) + estimator.set_input_shape('mask', (2, 1, x.size(2))) + estimator.set_input_shape('mu', (2, 80, x.size(2))) + estimator.set_input_shape('t', (2,)) + estimator.set_input_shape('spks', (2, 80)) + estimator.set_input_shape('cond', (2, 80, x.size(2))) + + data_ptrs = [x.contiguous().data_ptr(), + mask.contiguous().data_ptr(), + mu.contiguous().data_ptr(), + t.contiguous().data_ptr(), + spks.contiguous().data_ptr(), + cond.contiguous().data_ptr(), + x.data_ptr()] + + for idx, data_ptr in enumerate(data_ptrs): + estimator.set_tensor_address(engine.get_tensor_name(idx), data_ptr) + # run trt engine - self.estimator.execute_v2([x.contiguous().data_ptr(), - mask.contiguous().data_ptr(), - mu.contiguous().data_ptr(), - t.contiguous().data_ptr(), - spks.contiguous().data_ptr(), - cond.contiguous().data_ptr(), - x.data_ptr()]) - return x + estimator.execute_async_v3(torch.cuda.current_stream().cuda_stream) + + torch.cuda.current_stream().synchronize() + self.estimator.release_estimator(estimator) + return x + else: + with self.lock: + self.estimator.set_input_shape('x', (2, 80, x.size(2))) + self.estimator.set_input_shape('mask', (2, 1, x.size(2))) + self.estimator.set_input_shape('mu', (2, 80, x.size(2))) + self.estimator.set_input_shape('t', (2,)) + self.estimator.set_input_shape('spks', (2, 80)) + self.estimator.set_input_shape('cond', (2, 80, x.size(2))) + # run trt engine + self.estimator.execute_v2([x.contiguous().data_ptr(), + mask.contiguous().data_ptr(), + mu.contiguous().data_ptr(), + t.contiguous().data_ptr(), + spks.contiguous().data_ptr(), + cond.contiguous().data_ptr(), + x.data_ptr()]) + return x def compute_loss(self, x1, mask, mu, spks=None, cond=None): """Computes diffusion loss From fd9b7d45e2bac4c3c93077ba48d79890bb79e207 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=BE=E6=81=AF?= Date: Wed, 16 Apr 2025 11:24:51 +0800 Subject: [PATCH 2/4] Fix logging indentation in CosyVoice TTS method for improved clarity --- cosyvoice/cli/cosyvoice.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/cosyvoice/cli/cosyvoice.py b/cosyvoice/cli/cosyvoice.py index f49d2da..4c3b881 100644 --- a/cosyvoice/cli/cosyvoice.py +++ b/cosyvoice/cli/cosyvoice.py @@ -135,9 +135,9 @@ class CosyVoice: logging.info('synthesis text {}'.format(i)) for model_output in self.model.tts(**model_input, stream=stream, speed=speed): speech_len = model_output['tts_speech'].shape[1] / self.sample_rate - logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) - yield model_output - start_time = time.time() + logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) + yield model_output + start_time = time.time() cuda_stream.synchronize() self.stream_pool.put(cuda_stream) From 7f4c9a2c6454a1bb2a76e105069293a8370825c8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=BE=E6=81=AF?= Date: Wed, 16 Apr 2025 14:15:14 +0800 Subject: [PATCH 3/4] Refactor CosyVoice inference methods to streamline CUDA stream management - Removed the queue-based stream pool and integrated direct CUDA stream usage for improved performance. - Simplified inference methods by eliminating unnecessary synchronization and stream management code. - Enhanced logging for better tracking of synthesis operations and performance metrics. - Updated the model class to support CUDA stream context management, ensuring efficient resource utilization during inference. --- cosyvoice/cli/cosyvoice.py | 216 +++++++++++++------------------- cosyvoice/cli/model.py | 247 +++++++++++++++++++++---------------- 2 files changed, 226 insertions(+), 237 deletions(-) diff --git a/cosyvoice/cli/cosyvoice.py b/cosyvoice/cli/cosyvoice.py index 4c3b881..8606530 100644 --- a/cosyvoice/cli/cosyvoice.py +++ b/cosyvoice/cli/cosyvoice.py @@ -22,7 +22,7 @@ from cosyvoice.cli.frontend import CosyVoiceFrontEnd from cosyvoice.cli.model import CosyVoiceModel, CosyVoice2Model, VllmCosyVoice2Model from cosyvoice.utils.file_utils import logging from cosyvoice.utils.class_utils import get_model_type -import queue + class CosyVoice: @@ -60,11 +60,6 @@ class CosyVoice: self.fp16, self.estimator_count) del configs - thread_count = 10 - self.stream_pool = queue.Queue(maxsize=thread_count) - for _ in range(thread_count): - self.stream_pool.put(torch.cuda.Stream(self.device)) - def list_available_spks(self): spks = list(self.frontend.spk2info.keys()) @@ -74,104 +69,80 @@ class CosyVoice: self.frontend.add_spk_info(spk_id, spk_info) def inference_sft(self, tts_text, spk_id, stream=False, speed=1.0, text_frontend=True): - cuda_stream = self.stream_pool.get() - with torch.cuda.stream(cuda_stream): - for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): - model_input = self.frontend.frontend_sft(i, spk_id) - start_time = time.time() - logging.info('synthesis text {}'.format(i)) - for model_output in self.model.tts(**model_input, stream=stream, speed=speed): - speech_len = model_output['tts_speech'].shape[1] / self.sample_rate - logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) - yield model_output - start_time = time.time() - cuda_stream.synchronize() - self.stream_pool.put(cuda_stream) - - def inference_zero_shot(self, tts_text, prompt_text, prompt_speech_16k, stream=False, speed=1.0, text_frontend=True): - cuda_stream = self.stream_pool.get() - with torch.cuda.stream(cuda_stream): - prompt_text = self.frontend.text_normalize(prompt_text, split=False, text_frontend=text_frontend) - for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): - if (not isinstance(i, Generator)) and len(i) < 0.5 * len(prompt_text): - logging.warning('synthesis text {} too short than prompt text {}, this may lead to bad performance'.format(i, prompt_text)) - model_input = self.frontend.frontend_zero_shot(i, prompt_text, prompt_speech_16k, self.sample_rate) - start_time = time.time() - logging.info('synthesis text {}'.format(i)) - for model_output in self.model.tts(**model_input, stream=stream, speed=speed): - speech_len = model_output['tts_speech'].shape[1] / self.sample_rate - logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) - yield model_output - start_time = time.time() - cuda_stream.synchronize() - self.stream_pool.put(cuda_stream) - - def inference_zero_shot_by_spk_id(self, tts_text, spk_id, stream=False, speed=1.0, text_frontend=True): - """使用预定义的说话人执行 zero_shot 推理""" - cuda_stream = self.stream_pool.get() - with torch.cuda.stream(cuda_stream): - for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): - model_input = self.frontend.frontend_zero_shot_by_spk_id(i, spk_id) - start_time = time.time() - last_time = start_time - chunk_index = 0 - logging.info('synthesis text {}'.format(i)) - for model_output in self.model.tts(**model_input, stream=stream, speed=speed): - speech_len = model_output['tts_speech'].shape[1] / self.sample_rate - logging.info('yield speech index:{}, len {:.2f}, rtf {:.3f}, cost {:.3f}s, all cost time {:.3f}s'.format( - chunk_index, speech_len, (time.time()-last_time)/speech_len, time.time()-last_time, time.time()-start_time)) - yield model_output - last_time = time.time() - chunk_index += 1 - cuda_stream.synchronize() - self.stream_pool.put(cuda_stream) - - def inference_cross_lingual(self, tts_text, prompt_speech_16k, stream=False, speed=1.0, text_frontend=True): - cuda_stream = self.stream_pool.get() - with torch.cuda.stream(cuda_stream): - for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): - model_input = self.frontend.frontend_cross_lingual(i, prompt_speech_16k, self.sample_rate) - start_time = time.time() - logging.info('synthesis text {}'.format(i)) - for model_output in self.model.tts(**model_input, stream=stream, speed=speed): - speech_len = model_output['tts_speech'].shape[1] / self.sample_rate - logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) - yield model_output - start_time = time.time() - cuda_stream.synchronize() - self.stream_pool.put(cuda_stream) - - def inference_instruct(self, tts_text, spk_id, instruct_text, stream=False, speed=1.0, text_frontend=True): - cuda_stream = self.stream_pool.get() - with torch.cuda.stream(cuda_stream): - assert isinstance(self.model, CosyVoiceModel), 'inference_instruct is only implemented for CosyVoice!' - if self.instruct is False: - raise ValueError('{} do not support instruct inference'.format(self.model_dir)) - instruct_text = self.frontend.text_normalize(instruct_text, split=False, text_frontend=text_frontend) - for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): - model_input = self.frontend.frontend_instruct(i, spk_id, instruct_text) - start_time = time.time() - logging.info('synthesis text {}'.format(i)) - for model_output in self.model.tts(**model_input, stream=stream, speed=speed): - speech_len = model_output['tts_speech'].shape[1] / self.sample_rate - logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) - yield model_output - start_time = time.time() - cuda_stream.synchronize() - self.stream_pool.put(cuda_stream) - - def inference_vc(self, source_speech_16k, prompt_speech_16k, stream=False, speed=1.0): - cuda_stream = self.stream_pool.get() - with torch.cuda.stream(cuda_stream): - model_input = self.frontend.frontend_vc(source_speech_16k, prompt_speech_16k, self.sample_rate) + for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): + model_input = self.frontend.frontend_sft(i, spk_id) start_time = time.time() - for model_output in self.model.vc(**model_input, stream=stream, speed=speed): + logging.info('synthesis text {}'.format(i)) + for model_output in self.model.tts(**model_input, stream=stream, speed=speed): speech_len = model_output['tts_speech'].shape[1] / self.sample_rate logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) yield model_output start_time = time.time() - cuda_stream.synchronize() - self.stream_pool.put(cuda_stream) + + def inference_zero_shot(self, tts_text, prompt_text, prompt_speech_16k, stream=False, speed=1.0, text_frontend=True): + prompt_text = self.frontend.text_normalize(prompt_text, split=False, text_frontend=text_frontend) + for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): + if (not isinstance(i, Generator)) and len(i) < 0.5 * len(prompt_text): + logging.warning('synthesis text {} too short than prompt text {}, this may lead to bad performance'.format(i, prompt_text)) + model_input = self.frontend.frontend_zero_shot(i, prompt_text, prompt_speech_16k, self.sample_rate) + start_time = time.time() + logging.info('synthesis text {}'.format(i)) + for model_output in self.model.tts(**model_input, stream=stream, speed=speed): + speech_len = model_output['tts_speech'].shape[1] / self.sample_rate + logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) + yield model_output + start_time = time.time() + + def inference_zero_shot_by_spk_id(self, tts_text, spk_id, stream=False, speed=1.0, text_frontend=True): + """使用预定义的说话人执行 zero_shot 推理""" + for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): + model_input = self.frontend.frontend_zero_shot_by_spk_id(i, spk_id) + start_time = time.time() + last_time = start_time + chunk_index = 0 + logging.info('synthesis text {}'.format(i)) + for model_output in self.model.tts(**model_input, stream=stream, speed=speed): + speech_len = model_output['tts_speech'].shape[1] / self.sample_rate + logging.info('yield speech index:{}, len {:.2f}, rtf {:.3f}, cost {:.3f}s, all cost time {:.3f}s'.format( + chunk_index, speech_len, (time.time()-last_time)/speech_len, time.time()-last_time, time.time()-start_time)) + yield model_output + last_time = time.time() + chunk_index += 1 + + def inference_cross_lingual(self, tts_text, prompt_speech_16k, stream=False, speed=1.0, text_frontend=True): + for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): + model_input = self.frontend.frontend_cross_lingual(i, prompt_speech_16k, self.sample_rate) + start_time = time.time() + logging.info('synthesis text {}'.format(i)) + for model_output in self.model.tts(**model_input, stream=stream, speed=speed): + speech_len = model_output['tts_speech'].shape[1] / self.sample_rate + logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) + yield model_output + start_time = time.time() + + def inference_instruct(self, tts_text, spk_id, instruct_text, stream=False, speed=1.0, text_frontend=True): + assert isinstance(self.model, CosyVoiceModel), 'inference_instruct is only implemented for CosyVoice!' + if self.instruct is False: + raise ValueError('{} do not support instruct inference'.format(self.model_dir)) + instruct_text = self.frontend.text_normalize(instruct_text, split=False, text_frontend=text_frontend) + for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): + model_input = self.frontend.frontend_instruct(i, spk_id, instruct_text) + start_time = time.time() + logging.info('synthesis text {}'.format(i)) + for model_output in self.model.tts(**model_input, stream=stream, speed=speed): + speech_len = model_output['tts_speech'].shape[1] / self.sample_rate + logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) + yield model_output + start_time = time.time() + + def inference_vc(self, source_speech_16k, prompt_speech_16k, stream=False, speed=1.0): + model_input = self.frontend.frontend_vc(source_speech_16k, prompt_speech_16k, self.sample_rate) + start_time = time.time() + for model_output in self.model.vc(**model_input, stream=stream, speed=speed): + speech_len = model_output['tts_speech'].shape[1] / self.sample_rate + logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) + yield model_output + start_time = time.time() class CosyVoice2(CosyVoice): @@ -215,42 +186,29 @@ class CosyVoice2(CosyVoice): self.fp16, self.estimator_count) del configs - thread_count = 10 - self.stream_pool = queue.Queue(maxsize=thread_count) - for _ in range(thread_count): - self.stream_pool.put(torch.cuda.Stream(self.device)) def inference_instruct(self, *args, **kwargs): raise NotImplementedError('inference_instruct is not implemented for CosyVoice2!') def inference_instruct2(self, tts_text, instruct_text, prompt_speech_16k, stream=False, speed=1.0, text_frontend=True): - cuda_stream = self.stream_pool.get() - with torch.cuda.stream(cuda_stream): - assert isinstance(self.model, CosyVoice2Model), 'inference_instruct2 is only implemented for CosyVoice2!' - for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): - model_input = self.frontend.frontend_instruct2(i, instruct_text, prompt_speech_16k, self.sample_rate) + assert isinstance(self.model, CosyVoice2Model), 'inference_instruct2 is only implemented for CosyVoice2!' + for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): + model_input = self.frontend.frontend_instruct2(i, instruct_text, prompt_speech_16k, self.sample_rate) + start_time = time.time() + logging.info('synthesis text {}'.format(i)) + for model_output in self.model.tts(**model_input, stream=stream, speed=speed): + speech_len = model_output['tts_speech'].shape[1] / self.sample_rate + logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) + yield model_output start_time = time.time() - logging.info('synthesis text {}'.format(i)) - for model_output in self.model.tts(**model_input, stream=stream, speed=speed): - speech_len = model_output['tts_speech'].shape[1] / self.sample_rate - logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) - yield model_output - start_time = time.time() - cuda_stream.synchronize() - self.stream_pool.put(cuda_stream) def inference_instruct2_by_spk_id(self, tts_text, instruct_text, spk_id, stream=False, speed=1.0, text_frontend=True): - cuda_stream = self.stream_pool.get() - with torch.cuda.stream(cuda_stream): - assert isinstance(self.model, CosyVoice2Model), 'inference_instruct2 is only implemented for CosyVoice2!' - for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): - model_input = self.frontend.frontend_instruct2_by_spk_id(i, instruct_text, spk_id) + for i in tqdm(self.frontend.text_normalize(tts_text, split=True, text_frontend=text_frontend)): + model_input = self.frontend.frontend_instruct2_by_spk_id(i, instruct_text, spk_id) + start_time = time.time() + logging.info('synthesis text {}'.format(i)) + for model_output in self.model.tts(**model_input, stream=stream, speed=speed): + speech_len = model_output['tts_speech'].shape[1] / self.sample_rate + logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) + yield model_output start_time = time.time() - logging.info('synthesis text {}'.format(i)) - for model_output in self.model.tts(**model_input, stream=stream, speed=speed): - speech_len = model_output['tts_speech'].shape[1] / self.sample_rate - logging.info('yield speech len {}, rtf {}'.format(speech_len, (time.time() - start_time) / speech_len)) - yield model_output - start_time = time.time() - cuda_stream.synchronize() - self.stream_pool.put(cuda_stream) diff --git a/cosyvoice/cli/model.py b/cosyvoice/cli/model.py index 769dc92..d72816a 100644 --- a/cosyvoice/cli/model.py +++ b/cosyvoice/cli/model.py @@ -23,6 +23,7 @@ import uuid from cosyvoice.utils.common import fade_in_out from cosyvoice.utils.file_utils import convert_onnx_to_trt from cosyvoice.flow.flow_matching import EstimatorWrapper +import queue class CosyVoiceModel: @@ -66,6 +67,12 @@ class CosyVoiceModel: self.flow_cache_dict = {} self.hift_cache_dict = {} + self.stream_context_pool = queue.Queue() + for _ in range(10): + self.stream_context_pool.put(torch.cuda.stream(torch.cuda.Stream(self.device)) if torch.cuda.is_available() else nullcontext()) + + self.is_cuda_available = torch.cuda.is_available() + def load(self, llm_model, flow_model, hift_model): self.llm.load_state_dict(torch.load(llm_model, map_location=self.device), strict=True) self.llm.to(self.device).eval() @@ -166,63 +173,70 @@ class CosyVoiceModel: flow_prompt_speech_token=torch.zeros(1, 0, dtype=torch.int32), prompt_speech_feat=torch.zeros(1, 0, 80), stream=False, speed=1.0, **kwargs): # this_uuid is used to track variables related to this inference thread - this_uuid = str(uuid.uuid1()) - with self.lock: - self.tts_speech_token_dict[this_uuid], self.llm_end_dict[this_uuid] = [], False - self.hift_cache_dict[this_uuid] = None - self.mel_overlap_dict[this_uuid] = torch.zeros(1, 80, 0) - self.flow_cache_dict[this_uuid] = torch.zeros(1, 80, 0, 2) - p = threading.Thread(target=self.llm_job, args=(text, prompt_text, llm_prompt_speech_token, llm_embedding, this_uuid)) - p.start() - if stream is True: - token_hop_len = self.token_min_hop_len - while True: - time.sleep(0.1) - if len(self.tts_speech_token_dict[this_uuid]) >= token_hop_len + self.token_overlap_len: - this_tts_speech_token = torch.tensor(self.tts_speech_token_dict[this_uuid][:token_hop_len + self.token_overlap_len]) \ - .unsqueeze(dim=0) - this_tts_speech = self.token2wav(token=this_tts_speech_token, - prompt_token=flow_prompt_speech_token, - prompt_feat=prompt_speech_feat, - embedding=flow_embedding, - uuid=this_uuid, - finalize=False) - yield {'tts_speech': this_tts_speech.cpu()} - with self.lock: - self.tts_speech_token_dict[this_uuid] = self.tts_speech_token_dict[this_uuid][token_hop_len:] - # increase token_hop_len for better speech quality - token_hop_len = min(self.token_max_hop_len, int(token_hop_len * self.stream_scale_factor)) - if self.llm_end_dict[this_uuid] is True and len(self.tts_speech_token_dict[this_uuid]) < token_hop_len + self.token_overlap_len: - break - p.join() - # deal with remain tokens, make sure inference remain token len equals token_hop_len when cache_speech is not None - this_tts_speech_token = torch.tensor(self.tts_speech_token_dict[this_uuid]).unsqueeze(dim=0) - this_tts_speech = self.token2wav(token=this_tts_speech_token, - prompt_token=flow_prompt_speech_token, - prompt_feat=prompt_speech_feat, - embedding=flow_embedding, - uuid=this_uuid, - finalize=True) - yield {'tts_speech': this_tts_speech.cpu()} - else: - # deal with all tokens - p.join() - this_tts_speech_token = torch.tensor(self.tts_speech_token_dict[this_uuid]).unsqueeze(dim=0) - this_tts_speech = self.token2wav(token=this_tts_speech_token, - prompt_token=flow_prompt_speech_token, - prompt_feat=prompt_speech_feat, - embedding=flow_embedding, - uuid=this_uuid, - finalize=True, - speed=speed) - yield {'tts_speech': this_tts_speech.cpu()} - with self.lock: - self.tts_speech_token_dict.pop(this_uuid) - self.llm_end_dict.pop(this_uuid) - self.mel_overlap_dict.pop(this_uuid) - self.hift_cache_dict.pop(this_uuid) - self.flow_cache_dict.pop(this_uuid) - torch.cuda.empty_cache() + + stream_context = self.stream_context_pool.get() + with stream_context: + + this_uuid = str(uuid.uuid1()) + with self.lock: + self.tts_speech_token_dict[this_uuid], self.llm_end_dict[this_uuid] = [], False + self.hift_cache_dict[this_uuid] = None + self.mel_overlap_dict[this_uuid] = torch.zeros(1, 80, 0) + self.flow_cache_dict[this_uuid] = torch.zeros(1, 80, 0, 2) + p = threading.Thread(target=self.llm_job, args=(text, prompt_text, llm_prompt_speech_token, llm_embedding, this_uuid)) + p.start() + if stream is True: + token_hop_len = self.token_min_hop_len + while True: + time.sleep(0.1) + if len(self.tts_speech_token_dict[this_uuid]) >= token_hop_len + self.token_overlap_len: + this_tts_speech_token = torch.tensor(self.tts_speech_token_dict[this_uuid][:token_hop_len + self.token_overlap_len]) \ + .unsqueeze(dim=0) + this_tts_speech = self.token2wav(token=this_tts_speech_token, + prompt_token=flow_prompt_speech_token, + prompt_feat=prompt_speech_feat, + embedding=flow_embedding, + uuid=this_uuid, + finalize=False) + yield {'tts_speech': this_tts_speech.cpu()} + with self.lock: + self.tts_speech_token_dict[this_uuid] = self.tts_speech_token_dict[this_uuid][token_hop_len:] + # increase token_hop_len for better speech quality + token_hop_len = min(self.token_max_hop_len, int(token_hop_len * self.stream_scale_factor)) + if self.llm_end_dict[this_uuid] is True and len(self.tts_speech_token_dict[this_uuid]) < token_hop_len + self.token_overlap_len: + break + p.join() + # deal with remain tokens, make sure inference remain token len equals token_hop_len when cache_speech is not None + this_tts_speech_token = torch.tensor(self.tts_speech_token_dict[this_uuid]).unsqueeze(dim=0) + this_tts_speech = self.token2wav(token=this_tts_speech_token, + prompt_token=flow_prompt_speech_token, + prompt_feat=prompt_speech_feat, + embedding=flow_embedding, + uuid=this_uuid, + finalize=True) + yield {'tts_speech': this_tts_speech.cpu()} + else: + # deal with all tokens + p.join() + this_tts_speech_token = torch.tensor(self.tts_speech_token_dict[this_uuid]).unsqueeze(dim=0) + this_tts_speech = self.token2wav(token=this_tts_speech_token, + prompt_token=flow_prompt_speech_token, + prompt_feat=prompt_speech_feat, + embedding=flow_embedding, + uuid=this_uuid, + finalize=True, + speed=speed) + yield {'tts_speech': this_tts_speech.cpu()} + with self.lock: + self.tts_speech_token_dict.pop(this_uuid) + self.llm_end_dict.pop(this_uuid) + self.mel_overlap_dict.pop(this_uuid) + self.hift_cache_dict.pop(this_uuid) + self.flow_cache_dict.pop(this_uuid) + + self.synchronize_stream() + self.stream_context_pool.put(stream_context) + torch.cuda.empty_cache() def vc(self, source_speech_token, flow_prompt_speech_token, prompt_speech_feat, flow_embedding, stream=False, speed=1.0, **kwargs): # this_uuid is used to track variables related to this inference thread @@ -278,6 +292,10 @@ class CosyVoiceModel: self.hift_cache_dict.pop(this_uuid) torch.cuda.empty_cache() + def synchronize_stream(self): + if self.is_cuda_available: + torch.cuda.current_stream().synchronize() + class CosyVoice2Model(CosyVoiceModel): @@ -314,6 +332,12 @@ class CosyVoice2Model(CosyVoiceModel): self.llm_end_dict = {} self.hift_cache_dict = {} + self.stream_context_pool = queue.Queue() + for _ in range(10): + self.stream_context_pool.put(torch.cuda.stream(torch.cuda.Stream(self.device)) if torch.cuda.is_available() else nullcontext()) + + self.is_cuda_available = torch.cuda.is_available() + def load_jit(self, flow_encoder_model): flow_encoder = torch.jit.load(flow_encoder_model, map_location=self.device) self.flow.encoder = flow_encoder @@ -359,57 +383,64 @@ class CosyVoice2Model(CosyVoiceModel): flow_prompt_speech_token=torch.zeros(1, 0, dtype=torch.int32), prompt_speech_feat=torch.zeros(1, 0, 80), stream=False, speed=1.0, **kwargs): # this_uuid is used to track variables related to this inference thread - this_uuid = str(uuid.uuid1()) - with self.lock: - self.tts_speech_token_dict[this_uuid], self.llm_end_dict[this_uuid] = [], False - self.hift_cache_dict[this_uuid] = None - p = threading.Thread(target=self.llm_job, args=(text, prompt_text, llm_prompt_speech_token, llm_embedding, this_uuid)) - p.start() - if stream is True: - token_offset = 0 - while True: - time.sleep(0.1) - if len(self.tts_speech_token_dict[this_uuid]) - token_offset >= self.token_hop_len + self.flow.pre_lookahead_len: - this_tts_speech_token = torch.tensor(self.tts_speech_token_dict[this_uuid][:token_offset + self.token_hop_len + self.flow.pre_lookahead_len]).unsqueeze(dim=0) - this_tts_speech = self.token2wav(token=this_tts_speech_token, - prompt_token=flow_prompt_speech_token, - prompt_feat=prompt_speech_feat, - embedding=flow_embedding, - uuid=this_uuid, - token_offset=token_offset, - finalize=False) - token_offset += self.token_hop_len - yield {'tts_speech': this_tts_speech.cpu()} - if self.llm_end_dict[this_uuid] is True and len(self.tts_speech_token_dict[this_uuid]) - token_offset < self.token_hop_len + self.flow.pre_lookahead_len: - break - p.join() - # deal with remain tokens, make sure inference remain token len equals token_hop_len when cache_speech is not None - this_tts_speech_token = torch.tensor(self.tts_speech_token_dict[this_uuid]).unsqueeze(dim=0) - this_tts_speech = self.token2wav(token=this_tts_speech_token, - prompt_token=flow_prompt_speech_token, - prompt_feat=prompt_speech_feat, - embedding=flow_embedding, - uuid=this_uuid, - token_offset=token_offset, - finalize=True) - yield {'tts_speech': this_tts_speech.cpu()} - else: - # deal with all tokens - p.join() - this_tts_speech_token = torch.tensor(self.tts_speech_token_dict[this_uuid]).unsqueeze(dim=0) - this_tts_speech = self.token2wav(token=this_tts_speech_token, - prompt_token=flow_prompt_speech_token, - prompt_feat=prompt_speech_feat, - embedding=flow_embedding, - uuid=this_uuid, - token_offset=0, - finalize=True, - speed=speed) - yield {'tts_speech': this_tts_speech.cpu()} - with self.lock: - self.tts_speech_token_dict.pop(this_uuid) - self.llm_end_dict.pop(this_uuid) - torch.cuda.empty_cache() + self.synchronize_stream() + stream_context = self.stream_context_pool.get() + with torch.cuda.stream(stream_context): + + this_uuid = str(uuid.uuid1()) + with self.lock: + self.tts_speech_token_dict[this_uuid], self.llm_end_dict[this_uuid] = [], False + self.hift_cache_dict[this_uuid] = None + p = threading.Thread(target=self.llm_job, args=(text, prompt_text, llm_prompt_speech_token, llm_embedding, this_uuid)) + p.start() + if stream is True: + token_offset = 0 + while True: + time.sleep(0.1) + if len(self.tts_speech_token_dict[this_uuid]) - token_offset >= self.token_hop_len + self.flow.pre_lookahead_len: + this_tts_speech_token = torch.tensor(self.tts_speech_token_dict[this_uuid][:token_offset + self.token_hop_len + self.flow.pre_lookahead_len]).unsqueeze(dim=0) + this_tts_speech = self.token2wav(token=this_tts_speech_token, + prompt_token=flow_prompt_speech_token, + prompt_feat=prompt_speech_feat, + embedding=flow_embedding, + uuid=this_uuid, + token_offset=token_offset, + finalize=False) + token_offset += self.token_hop_len + yield {'tts_speech': this_tts_speech.cpu()} + if self.llm_end_dict[this_uuid] is True and len(self.tts_speech_token_dict[this_uuid]) - token_offset < self.token_hop_len + self.flow.pre_lookahead_len: + break + p.join() + # deal with remain tokens, make sure inference remain token len equals token_hop_len when cache_speech is not None + this_tts_speech_token = torch.tensor(self.tts_speech_token_dict[this_uuid]).unsqueeze(dim=0) + this_tts_speech = self.token2wav(token=this_tts_speech_token, + prompt_token=flow_prompt_speech_token, + prompt_feat=prompt_speech_feat, + embedding=flow_embedding, + uuid=this_uuid, + token_offset=token_offset, + finalize=True) + yield {'tts_speech': this_tts_speech.cpu()} + else: + # deal with all tokens + p.join() + this_tts_speech_token = torch.tensor(self.tts_speech_token_dict[this_uuid]).unsqueeze(dim=0) + this_tts_speech = self.token2wav(token=this_tts_speech_token, + prompt_token=flow_prompt_speech_token, + prompt_feat=prompt_speech_feat, + embedding=flow_embedding, + uuid=this_uuid, + token_offset=0, + finalize=True, + speed=speed) + yield {'tts_speech': this_tts_speech.cpu()} + with self.lock: + self.tts_speech_token_dict.pop(this_uuid) + self.llm_end_dict.pop(this_uuid) + + self.synchronize_stream() + self.stream_context_pool.put(stream_context) + torch.cuda.empty_cache() class VllmCosyVoice2Model(CosyVoice2Model): From 369f3c2c18c68867a06cd4c1dd8d58acd619a1b9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=A6=BE=E6=81=AF?= Date: Wed, 16 Apr 2025 14:39:06 +0800 Subject: [PATCH 4/4] Update estimator count retrieval and memory pool limit in CosyVoice - Simplified estimator count retrieval in CosyVoice and CosyVoice2 classes to directly access the configs dictionary. - Adjusted memory pool limit in the ONNX to TensorRT conversion function from 8GB to 1GB for optimized resource management. --- cosyvoice/cli/cosyvoice.py | 4 ++-- cosyvoice/utils/file_utils.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/cosyvoice/cli/cosyvoice.py b/cosyvoice/cli/cosyvoice.py index 8606530..7f0211d 100644 --- a/cosyvoice/cli/cosyvoice.py +++ b/cosyvoice/cli/cosyvoice.py @@ -54,7 +54,7 @@ class CosyVoice: '{}/llm.llm.{}.zip'.format(model_dir, 'fp16' if self.fp16 is True else 'fp32'), '{}/flow.encoder.{}.zip'.format(model_dir, 'fp16' if self.fp16 is True else 'fp32')) if load_trt: - self.estimator_count = configs['flow']['decoder']['estimator'].get('estimator_count', 1) + self.estimator_count = configs.get('estimator_count', 1) self.model.load_trt('{}/flow.decoder.estimator.{}.mygpu.plan'.format(model_dir, 'fp16' if self.fp16 is True else 'fp32'), '{}/flow.decoder.estimator.fp32.onnx'.format(model_dir), self.fp16, self.estimator_count) @@ -180,7 +180,7 @@ class CosyVoice2(CosyVoice): if load_jit: self.model.load_jit('{}/flow.encoder.{}.zip'.format(model_dir, 'fp16' if self.fp16 is True else 'fp32')) if load_trt: - self.estimator_count = configs['flow']['decoder']['estimator'].get('estimator_count', 1) + self.estimator_count = configs.get('estimator_count', 1) self.model.load_trt('{}/flow.decoder.estimator.{}.mygpu.plan'.format(model_dir, 'fp16' if self.fp16 is True else 'fp32'), '{}/flow.decoder.estimator.fp32.onnx'.format(model_dir), self.fp16, self.estimator_count) diff --git a/cosyvoice/utils/file_utils.py b/cosyvoice/utils/file_utils.py index ac7fe93..cf8ad03 100644 --- a/cosyvoice/utils/file_utils.py +++ b/cosyvoice/utils/file_utils.py @@ -61,7 +61,7 @@ def convert_onnx_to_trt(trt_model, onnx_model, fp16): network = builder.create_network(network_flags) parser = trt.OnnxParser(network, logger) config = builder.create_builder_config() - config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 33) # 8GB + config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30) # 1GB if fp16: config.set_flag(trt.BuilderFlag.FP16) profile = builder.create_optimization_profile()