diff --git a/src/engine_args.py b/src/engine_args.py index 45e50d1..3131172 100644 --- a/src/engine_args.py +++ b/src/engine_args.py @@ -4,6 +4,7 @@ import logging from torch.cuda import device_count from vllm import AsyncEngineArgs from vllm.model_executor.model_loader.tensorizer import TensorizerConfig +from src.utils import convert_limit_mm_per_prompt RENAME_ARGS_MAP = { "MODEL_NAME": "model", @@ -89,7 +90,8 @@ DEFAULT_ARGS = { "qlora_adapter_name_or_path": os.getenv('QLORA_ADAPTER_NAME_OR_PATH', None), "disable_logprobs_during_spec_decoding": os.getenv('DISABLE_LOGPROBS_DURING_SPEC_DECODING', None), "otlp_traces_endpoint": os.getenv('OTLP_TRACES_ENDPOINT', None), - "use_v2_block_manager": os.getenv('USE_V2_BLOCK_MANAGER', 'true') + "use_v2_block_manager": os.getenv('USE_V2_BLOCK_MANAGER', 'true'), + "limit_mm_per_prompt": convert_limit_mm_per_prompt(os.getenv('LIMIT_MM_PER_PROMPT', "image=1")) } def match_vllm_args(args): diff --git a/src/utils.py b/src/utils.py index bfc8ce9..85af1e8 100644 --- a/src/utils.py +++ b/src/utils.py @@ -15,6 +15,10 @@ except ImportError: logging.basicConfig(level=logging.INFO) +def convert_limit_mm_per_prompt(input_string: str): + key, value = input_string.split('=') + return {key: int(value)} + def count_physical_cores(): with open('/proc/cpuinfo') as f: content = f.readlines()