2121from vllm .multimodal .image import convert_image_mode
2222from vllm .utils .argparse_utils import FlexibleArgumentParser
2323
24+ # Import StageConfigFactory for Tier-2 CLI override testing
2425from vllm_omni .entrypoints .omni import Omni
2526
2627SEED = 42
@@ -325,11 +326,30 @@ def main(args):
325326 else :
326327 query_result = query_func ()
327328
329+ # Build kwargs with Tier-2 CLI overrides
330+ omni_kwargs = {
331+ "stage_configs_path" : args .stage_configs_path ,
332+ "log_stats" : args .log_stats ,
333+ "stage_init_timeout" : args .stage_init_timeout ,
334+ }
335+
336+ # Add Tier-2 CLI overrides if specified
337+ if args .gpu_memory_utilization is not None :
338+ omni_kwargs ["gpu_memory_utilization" ] = args .gpu_memory_utilization
339+ if args .tensor_parallel_size is not None :
340+ omni_kwargs ["tensor_parallel_size" ] = args .tensor_parallel_size
341+ if args .devices is not None :
342+ omni_kwargs ["devices" ] = args .devices
343+ if args .enforce_eager :
344+ omni_kwargs ["enforce_eager" ] = args .enforce_eager
345+ if args .trust_remote_code :
346+ omni_kwargs ["trust_remote_code" ] = args .trust_remote_code
347+ if args .stage_id is not None :
348+ omni_kwargs ["stage_id" ] = args .stage_id
349+
328350 omni_llm = Omni (
329351 model = model_name ,
330- stage_configs_path = args .stage_configs_path ,
331- log_stats = args .log_stats ,
332- stage_init_timeout = args .stage_init_timeout ,
352+ ** omni_kwargs ,
333353 )
334354
335355 thinker_sampling_params = SamplingParams (
@@ -489,6 +509,12 @@ def parse_args():
489509 default = "output_audio" ,
490510 help = "[Deprecated] Output wav directory (use --output-dir)." ,
491511 )
512+ parser .add_argument (
513+ "--output-dir" ,
514+ type = str ,
515+ default = None ,
516+ help = "Output directory for generated files (text and audio)." ,
517+ )
492518 parser .add_argument (
493519 "--num-prompts" ,
494520 type = int ,
@@ -505,7 +531,44 @@ def parse_args():
505531 "--stage-configs-path" ,
506532 type = str ,
507533 default = None ,
508- help = "Path to a stage configs file." ,
534+ help = "Path to a stage configs file. If not specified, uses auto-detected Tier-1 topology." ,
535+ )
536+ # Tier-2 CLI override arguments
537+ parser .add_argument (
538+ "--gpu-memory-utilization" ,
539+ type = float ,
540+ default = None ,
541+ help = "GPU memory utilization for all stages (Tier-2 override). Example: 0.9" ,
542+ )
543+ parser .add_argument (
544+ "--tensor-parallel-size" ,
545+ type = int ,
546+ default = None ,
547+ help = "Tensor parallel size for all stages (Tier-2 override). Example: 2" ,
548+ )
549+ parser .add_argument (
550+ "--devices" ,
551+ type = str ,
552+ default = None ,
553+ help = "Device assignment for stages (Tier-2 override). Example: '0,1'" ,
554+ )
555+ parser .add_argument (
556+ "--enforce-eager" ,
557+ action = "store_true" ,
558+ default = False ,
559+ help = "Enforce eager mode for all stages (Tier-2 override)." ,
560+ )
561+ parser .add_argument (
562+ "--trust-remote-code" ,
563+ action = "store_true" ,
564+ default = False ,
565+ help = "Trust remote code for model loading (Tier-2 override)." ,
566+ )
567+ parser .add_argument (
568+ "--stage-id" ,
569+ type = int ,
570+ default = None ,
571+ help = "Launch only the specified stage ID for independent stage testing." ,
509572 )
510573 parser .add_argument (
511574 "--video-path" ,
0 commit comments