[update] support CosyVoice2

api_server.py

@@ -244,7 +244,7 @@ async def create_speech(
     try:
         request_data = await request.json()
         model = request_data.get("model")
-        voice = request_data.get("voice", "alloy")
+        voice = request_data.get("voice", "prompt_data")
         response_format = request_data.get("response_format", "mp3")
 
         if not model:

backend/tts_client_backend.py

@@ -11,6 +11,7 @@ from .base_model_backend import BaseModelBackend
 from client.tts_client import TTSClient
 from concurrent.futures import ThreadPoolExecutor
 from services.memory_check import MemoryChecker
+import tiktoken
 
 class TtsClientBackend(BaseModelBackend):
     SUPPORTED_FORMATS = ["mp3", "opus", "aac", "flac", "wav", "pcm"]
@@ -21,6 +22,7 @@ class TtsClientBackend(BaseModelBackend):
         self._active_clients = {}
         self._pool_lock = asyncio.Lock()
         self.logger = logging.getLogger("api.tts")
+        self.MAX_CONTEXT_LENGTH = model_config.get("max_context_length", 256)
         self.POOL_SIZE = 1
         self._inference_executor = ThreadPoolExecutor(max_workers=self.POOL_SIZE)
         self._active_tasks = weakref.WeakSet()
@@ -28,10 +30,11 @@ class TtsClientBackend(BaseModelBackend):
             host=self.config["host"],
             port=self.config["port"]
         )
-        self.sample_rate = 16000
+        self.sample_rate =  model_config.get("sample_rate", 16000)
         self.channels = 1
+        self.tokenizer = tiktoken.get_encoding("cl100k_base")
 
-    async def _get_client(self):
+    async def _get_client(self, voice: str = "prompt_data"):
         try:
             await asyncio.wait_for(self._pool_lock.acquire(), timeout=30.0)
             
@@ -61,13 +64,13 @@ class TtsClientBackend(BaseModelBackend):
                 await loop.run_in_executor(
                     self._inference_executor,
                     lambda: client.setup(
-                        "melotts.setup",
+                        self.config["object"],
                         {
                             "model": self.config["model_name"],
                             "response_format": "pcm.stream.base64",
                             "input": "tts.utf-8",
                             "enoutput": True,
-                            "voice": "alloy"
+                            "prompt_dir": voice
                         }
                     )
                 )
@@ -129,9 +132,21 @@ class TtsClientBackend(BaseModelBackend):
         self._full_audio_buffer.write(pcm_data)
         
         return b''
+    
+    def _count_tokens(self, text: str) -> int:
+        """Count the number of tokens in a given text."""
+        return len(self.tokenizer.encode(text))
 
-    async def generate_speech(self, input_text: str, voice: str = "alloy", format: str = "mp3") -> AsyncGenerator[bytes, None]:
-        client = await self._get_client()
+    async def generate_speech(self, input_text: str, voice: str = "prompt_data", format: str = "mp3") -> AsyncGenerator[bytes, None]:
+        token_count = self._count_tokens(input_text)
+        if token_count > self.MAX_CONTEXT_LENGTH:
+            msg = (
+                f"Input token count ({token_count}) exceeds max context length ({self.MAX_CONTEXT_LENGTH})."
+            )
+            self.logger.warning(msg)
+            raise ValueError(msg)
+
+        client = await self._get_client(voice)
         task = asyncio.current_task()
         self._active_tasks.add(task)
         full_data = b''

services/model_list.py

@@ -94,23 +94,65 @@ class GetModelList:
                             new_entry['max_context_length'] = precompute_len
 
                     elif model_type == 'tts':
+                        mode_param = model_data.get("mode_param", {})
+                        precompute_len = None
+                        sample_rate = None
+                        cmm_size = None
+                        if isinstance(mode_param, dict):
+                            precompute_len = mode_param.get("precompute_len")
+                            cmm_size = mode_param.get("cmm_size")
+                            sample_rate = mode_param.get("sample_rate")
+
                         if 'melotts' in mode.lower():
                             obj = 'melotts.setup'
                             new_entry['memory_required'] = 59764
+                            new_entry['sample_rate'] = 16000
+                        elif 'cosyvoice' in mode.lower():
+                            obj = 'cosy_voice.setup'
+                            new_entry['memory_required'] = 1185772
+                            new_entry['sample_rate'] = 48000
                         else:
                             obj = 'tts.setup'
-        
+
+                        if cmm_size is not None:
+                            new_entry['memory_required'] = cmm_size
+                        if precompute_len is not None:
+                            new_entry['max_context_length'] = precompute_len
+                        if sample_rate is not None:
+                            new_entry['max_context_length'] = sample_rate
+
                         new_entry.update({
                             "response_format": "wav.base64",
                             "object": obj
                         })
                     elif model_type == 'asr':
+                        mode_param = model_data.get("mode_param", {})
+                        precompute_len = None
+                        cmm_size = None
+                        if isinstance(mode_param, dict):
+                            precompute_len = mode_param.get("precompute_len")
+                            cmm_size = mode_param.get("cmm_size")
+
                         if 'whisper' in mode.lower():
                             obj = 'whisper.setup'
                             if 'tiny' in mode:
-                                new_entry['memory_required'] = 289132
+                                new_entry['memory_required'] = 263860
+                            elif 'base' in mode.lower():
+                                new_entry['memory_required'] = 448212
+                            elif 'small' in mode.lower():
+                                new_entry['memory_required'] = 1132748
+                            elif 'turbo' in mode.lower():
+                                new_entry['memory_required'] = 2048000
+                            else:
+                                new_entry['memory_required'] = 1500000
                         else:
                             obj = 'asr.setup'
+
+                        if cmm_size is not None:
+                            new_entry['memory_required'] = cmm_size
+                        if precompute_len is not None:
+                            new_entry['max_context_length'] = precompute_len
+
                         new_entry.update({
                             "input": "pcm.base64",
                             "response_format": "asr.utf-8",
@@ -133,4 +175,4 @@ class GetModelList:
         return await loop.run_in_executor(
             None, 
             self._sys_client.model_list
-        ) 
+        )