vllm-omni/vllm_omni/entrypoints/cli/main.py

"""
CLI entry point for vLLM-omni that intercepts vLLM commands.
"""

import sys

from vllm_omni.entrypoints.cli.serve import OmniServeCommand


def main():
    """Main CLI entry point that intercepts vLLM commands."""
    # Check if --omni flag is present
    if "--omni" in sys.argv:
        # Remove --omni flag and "serve" command, process with vLLM-omni
        omni_args = [arg for arg in sys.argv[1:] if arg not in ["--omni", "serve"]]
        omni_serve = OmniServeCommand()
        omni_serve.run(omni_args)
    else:
        # Forward to original vLLM CLI
        from vllm.entrypoints.cli.main import main as vllm_main

        vllm_main()


if __name__ == "__main__":
    main()