diff --git a/backend/llm_client_backend.py b/backend/llm_client_backend.py index 635398f..efb1e1a 100644 --- a/backend/llm_client_backend.py +++ b/backend/llm_client_backend.py @@ -77,8 +77,8 @@ class LlmClientBackend(BaseModelBackend): await asyncio.sleep(retry_interval) await asyncio.wait_for(self._pool_lock.acquire(), timeout=timeout - (time.time() - start_time)) - # if "memory_required" in self.config: - # await self.memory_checker.check_memory(self.config["memory_required"]) + if "memory_required" in self.config: + await self.memory_checker.check_memory(self.config["memory_required"]) self.logger.debug("Creating new LLM client") client = LLMClient( diff --git a/services/model_list.py b/services/model_list.py index b27979f..85e65e8 100644 --- a/services/model_list.py +++ b/services/model_list.py @@ -56,6 +56,13 @@ class GetModelList: "object": f"{model_type}.setup", "system_prompt": "You are a helpful assistant." }) + + 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 '-1.5B-' in mode: new_entry['memory_required'] = 1782579 new_entry['pool_size'] = 1 @@ -68,6 +75,8 @@ class GetModelList: else: new_entry['memory_required'] = 1363148 new_entry['pool_size'] = 2 + if cmm_size is not None: + new_entry['memory_required'] = cmm_size if '-p256-' in mode: new_entry['max_context_length'] = 256 @@ -81,6 +90,9 @@ class GetModelList: if '-ctx-' in mode: new_entry['max_context_length'] = 1024 + if precompute_len is not None: + new_entry['max_context_length'] = precompute_len + elif model_type == 'tts': if 'melotts' in mode.lower(): obj = 'melotts.setup'