diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index be1a1bf14..3b222c817 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -33,7 +33,7 @@ Examples: vllm serve Qwen/Qwen-Image --omni --port 8091 Search by using: `--help=` to explore options by section (e.g., ---help=ModelConfig, --help=Frontend) +--help=OmniConfig) Use `--help=all` to show all available flags at once. """ @@ -70,60 +70,67 @@ class OmniServeCommand(CLISubcommand): serve_parser = make_arg_parser(serve_parser) serve_parser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(subcmd=self.name) - serve_parser.add_argument( + + # Create OmniConfig argument group for omni-related parameters + # This ensures the parameters appear in --help output + omni_config_group = serve_parser.add_argument_group( + title="OmniConfig", description="Configuration for vLLM-Omni multi-stage and diffusion models." + ) + + omni_config_group.add_argument( "--omni", action="store_true", help="Enable vLLM-Omni mode for multi-modal and diffusion models", ) - serve_parser.add_argument( + omni_config_group.add_argument( "--stage-configs-path", type=str, default=None, help="Path to the stage configs file. If not specified, the stage configs will be loaded from the model.", ) - serve_parser.add_argument( + omni_config_group.add_argument( "--stage-init-timeout", type=int, default=300, help="The timeout for initializing a single stage in seconds (default: 300)", ) - serve_parser.add_argument( + omni_config_group.add_argument( "--init-timeout", type=int, default=60000, help="The timeout for initializing the stages.", ) - serve_parser.add_argument( + omni_config_group.add_argument( "--shm-threshold-bytes", type=int, default=65536, help="The threshold for the shared memory size.", ) - serve_parser.add_argument( + omni_config_group.add_argument( "--log-stats", action="store_true", help="Enable logging the stats.", ) - serve_parser.add_argument( + omni_config_group.add_argument( "--log-file", type=str, default=None, help="The path to the log file.", ) - serve_parser.add_argument( + omni_config_group.add_argument( "--batch-timeout", type=int, default=10, help="The timeout for the batch.", ) - serve_parser.add_argument( + omni_config_group.add_argument( "--worker-backend", type=str, default="multi_process", choices=["multi_process", "ray"], help="The backend to use for stage workers.", ) - serve_parser.add_argument( + omni_config_group.add_argument( "--ray-address", type=str, default=None, @@ -131,13 +138,13 @@ class OmniServeCommand(CLISubcommand): ) # Diffusion model specific arguments - serve_parser.add_argument( + omni_config_group.add_argument( "--num-gpus", type=int, default=None, help="Number of GPUs to use for diffusion model inference.", ) - serve_parser.add_argument( + omni_config_group.add_argument( "--usp", "--ulysses-degree", dest="ulysses_degree", @@ -146,7 +153,7 @@ class OmniServeCommand(CLISubcommand): help="Ulysses Sequence Parallelism degree for diffusion models. " "Equivalent to setting DiffusionParallelConfig.ulysses_degree.", ) - serve_parser.add_argument( + omni_config_group.add_argument( "--ring", dest="ring_degree", type=int, @@ -156,13 +163,13 @@ class OmniServeCommand(CLISubcommand): ) # Cache optimization parameters - serve_parser.add_argument( + omni_config_group.add_argument( "--cache-backend", type=str, default="none", help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit'", ) - serve_parser.add_argument( + omni_config_group.add_argument( "--cache-config", type=str, default=None, @@ -170,12 +177,12 @@ class OmniServeCommand(CLISubcommand): ) # VAE memory optimization parameters - serve_parser.add_argument( + omni_config_group.add_argument( "--vae-use-slicing", action="store_true", help="Enable VAE slicing for memory optimization (useful for mitigating OOM issues).", ) - serve_parser.add_argument( + omni_config_group.add_argument( "--vae-use-tiling", action="store_true", help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).", @@ -189,13 +196,13 @@ class OmniServeCommand(CLISubcommand): ) # Video model parameters (e.g., Wan2.2) - engine-level - serve_parser.add_argument( + omni_config_group.add_argument( "--boundary-ratio", type=float, default=None, help="Boundary split ratio for low/high DiT in video models (e.g., 0.875 for Wan2.2).", ) - serve_parser.add_argument( + omni_config_group.add_argument( "--flow-shift", type=float, default=None,