diff --git a/api_server.py b/api_server.py index 95f96d8..7a82136 100644 --- a/api_server.py +++ b/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: diff --git a/backend/tts_client_backend.py b/backend/tts_client_backend.py index 8039ca6..34e4db9 100644 --- a/backend/tts_client_backend.py +++ b/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,22 @@ 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]: + if self._count_tokens(input_text) > self.MAX_CONTEXT_LENGTH: + self.logger.warning( + f"Input text length ({len(input_text)}) exceeds max context length ({self.MAX_CONTEXT_LENGTH})." + ) + raise ValueError( + f"Text too long: {len(input_text)} characters. " + f"Maximum allowed: {self.MAX_CONTEXT_LENGTH}." + ) + + client = await self._get_client(voice) task = asyncio.current_task() self._active_tasks.add(task) full_data = b'' diff --git a/services/model_list.py b/services/model_list.py index 85e65e8..96178be 100644 --- a/services/model_list.py +++ 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 - ) + ) \ No newline at end of file