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

@@ -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,8 +69,6 @@ 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()
@@ -85,12 +78,8 @@ class CosyVoice:
                 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):
@@ -103,13 +92,9 @@ class CosyVoice:
                 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()
@@ -123,12 +108,8 @@ class CosyVoice:
                 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()
@@ -138,12 +119,8 @@ class CosyVoice:
                 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))
@@ -157,12 +134,8 @@ class CosyVoice:
                 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()
         for model_output in self.model.vc(**model_input, stream=stream, speed=speed):
@@ -170,8 +143,6 @@ class CosyVoice:
             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):
@@ -215,17 +186,11 @@ 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)
@@ -236,13 +201,8 @@ class CosyVoice2(CosyVoice):
                 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)
             start_time = time.time()
@@ -252,5 +212,3 @@ class CosyVoice2(CosyVoice):
                 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)

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,6 +173,10 @@ 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
+
+        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
@@ -222,6 +233,9 @@ class CosyVoiceModel:
                 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):
@@ -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,6 +383,10 @@ 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
+        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
@@ -409,6 +437,9 @@ class CosyVoice2Model(CosyVoiceModel):
             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()