[Feat] Rename AsyncOmniLLM -> AsyncOmni (#103)

Signed-off-by: WANG Cong <115451386+congw729@users.noreply.github.com>

docs/api/README.md

@@ -4,8 +4,8 @@
 
 Main entry points for vLLM-Omni inference and serving.
 
-- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
-- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniStageLLM][]
+- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
+- [vllm_omni.entrypoints.async_omni.AsyncOmniStageLLM][]
 - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalContentParser][]
 - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalItemTracker][]
 - [vllm_omni.entrypoints.chat_utils.parse_chat_messages_futures][]
@@ -102,10 +102,6 @@ Configuration classes.
 
 Worker classes and model runners for distributed inference.
 
-- [vllm_omni.diffusion.worker.gpu_worker.GPUWorker][]
-- [vllm_omni.diffusion.worker.gpu_worker.WorkerProc][]
-- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorker][]
-- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorkerProc][]
 - [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][]
 - [vllm_omni.worker.gpu_ar_worker.GPUARWorker][]
 - [vllm_omni.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][]

docs/configuration/stage_configs.md

@@ -44,7 +44,7 @@ stage_args:
       trust_remote_code: true # Needed by huggingface config parsing
       engine_output_type: latent  # It claims that the stage will input latent hiddenstates besides token ids
       enable_prefix_caching: false # For request with hiddenstates output, the prefix caching is not supported now
-    is_comprehension: true # If the stage is a text or multimodal comprehension module. If it is, the AsyncOmniLLM will use its tokenizer as default
+    is_comprehension: true # If the stage is a text or multimodal comprehension module. If it is, the AsyncOmni will use its tokenizer as default
     final_output: true # If the stage has output as part of final outputs. If it is false, which means that the stage only works as a intermediate role.
     final_output_type: text # What is the final output type. It can be text and audio now.
     default_sampling_params: # sampling parameters for the stage. Their meaning aligns with vLLM.
@@ -206,7 +206,7 @@ Default: `false`
 
 ### `is_comprehension`
 
-Whether this stage is a text or multimodal comprehension module. When set to `true`, the stage acts as a comprehension module that processes input text or multimodal content. If this is the first comprehension stage, `AsyncOmniLLM` will use its tokenizer as the default tokenizer for the entire pipeline.
+Whether this stage is a text or multimodal comprehension module. When set to `true`, the stage acts as a comprehension module that processes input text or multimodal content. If this is the first comprehension stage, `AsyncOmni` will use its tokenizer as the default tokenizer for the entire pipeline.
 
 Default: `true`
 

docs/design/vllm_omni_design.md

@@ -14,7 +14,7 @@ vLLM-Omni is a multi-modality extension for vLLM that supports non-autoregressiv
 
 ## Key Data Flow
 
-API Server --> OmniLLM/AsyncOmniLLM (New, including multi engines) --> LLMEngine/AsyncLLM --> Engine Core
+API Server --> OmniLLM/AsyncOmni (New, including multi engines) --> LLMEngine/AsyncLLM --> Engine Core
  --> Scheduler (New one for DiT) --> Executor (New one for diffusers) --> Worker (New one for DiT)
  --> ModelRunner (New one for AR hiddenstate, New one for DiT) --> RequestState --> OutputProcessoer (New one for final multimodal output)
 
@@ -40,7 +40,7 @@ graph TD
     B --> C{Detect --omni flag}
     C -->|Yes| D[Parse OmniConfig]
     C -->|No| E[Forward to vLLM CLI]
-    D --> F[Initialize AsyncOmniLLM]
+    D --> F[Initialize AsyncOmni]
     F --> G[Start omni Server]
     G --> H[Multi-stage Processing]
     E --> I[Standard vLLM Pipeline]
@@ -209,7 +209,7 @@ Similar to OmniLLM in offline inference, add some asynchronous processing, refer
 from vllm.v1.engine.async_llm import AsyncLLM
 
 
-class AsyncOmniLLM(AsyncLLM):
+class AsyncOmni(AsyncLLM):
     """Extended AsyncLLM supporting multiple engines and stage-based processing"""
 
     def __init__(self, stage_configs: List[StageConfig]):

docs/mkdocs/hooks/generate_api_readme.py

@@ -188,6 +188,11 @@ def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]:
             # Add classes (filter out internal ones)
             for class_name in classes:
                 class_short_name = class_name.split(".")[-1]
+
+                # TODO: Implement worker API reference
+                if "diffusion.worker.gpu_worker" in class_name:
+                    continue
+
                 # Skip if it matches internal patterns (unless it's a main model class)
                 if any(pattern in class_short_name for pattern in internal_patterns):
                     # But include main model classes

docs/user_guide/examples/online_serving/qwen2_5_omni.md

@@ -48,7 +48,7 @@ This Web UI demo allows users to interact with the model through a web browser.
 
 ### Running Gradio Demo
 
-Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by
 
 ```bash
 python gradio_demo.py  --model Qwen/Qwen2.5-Omni-7B --port 7861

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -48,7 +48,7 @@ This Web UI demo allows users to interact with the model through a web browser.
 
 ### Running Gradio Demo
 
-Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by
 
 ```bash
 python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861
@@ -72,7 +72,7 @@ python gradio_demo.py \
 ```
 
 - `--model`: Model name
-- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmniLLM locally.
+- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmni locally.
 - `--api-base`: Base URL for vllm serve (only used when `use-api-server` is set, default: http://localhost:8091/v1)
 - `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
 - `--port`: Port for Gradio server (default: 7861)

examples/online_serving/qwen2_5_omni/README.md

@@ -45,7 +45,7 @@ This Web UI demo allows users to interact with the model through a web browser.
 
 ### Running Gradio Demo
 
-Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by
 
 ```bash
 python gradio_demo.py  --model Qwen/Qwen2.5-Omni-7B --port 7861

examples/online_serving/qwen2_5_omni/gradio_demo.py

@@ -15,7 +15,7 @@ from PIL import Image
 from vllm.assets.video import video_get_metadata, video_to_ndarrays
 from vllm.sampling_params import SamplingParams
 
-from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
+from vllm_omni.entrypoints.async_omni import AsyncOmni
 
 # Import utils from offline inference example
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../offline_inference/qwen2_5_omni"))
@@ -236,8 +236,8 @@ def process_video_file(
     return frames, metadata, audio_tuple
 
 
-async def run_inference_async_omni_llm(
-    omni_llm: AsyncOmniLLM,
+async def run_inference_async_omni(
+    omni: AsyncOmni,
     sampling_params: list[SamplingParams],
     prompt_args_template: SimpleNamespace,
     user_prompt: str,
@@ -246,7 +246,7 @@ async def run_inference_async_omni_llm(
     video_file: Optional[str] = None,
     use_audio_in_video: bool = False,
 ):
-    """Run inference using AsyncOmniLLM directly with multimodal support."""
+    """Run inference using AsyncOmni directly with multimodal support."""
     if not user_prompt.strip() and not audio_file and not image_file and not video_file:
         return "Please provide at least a text prompt or multimodal input.", None
 
@@ -328,7 +328,7 @@ async def run_inference_async_omni_llm(
         text_outputs: list[str] = []
         audio_output = None
 
-        async for stage_outputs in omni_llm.generate(
+        async for stage_outputs in omni.generate(
             prompt=omni_prompt,
             request_id=request_id,
             sampling_params_list=sampling_params,
@@ -362,12 +362,12 @@ async def run_inference_async_omni_llm(
 
 
 def build_interface(
-    omni_llm: AsyncOmniLLM,
+    omni: AsyncOmni,
     sampling_params: list[SamplingParams],
     prompt_args_template: SimpleNamespace,
     model: str,
 ):
-    """Build Gradio interface for AsyncOmniLLM mode."""
+    """Build Gradio interface for AsyncOmni mode."""
 
     async def run_inference(
         user_prompt: str,
@@ -376,8 +376,8 @@ def build_interface(
         video_file: Optional[str],
         use_audio_in_video: bool,
     ):
-        return await run_inference_async_omni_llm(
-            omni_llm,
+        return await run_inference_async_omni(
+            omni,
             sampling_params,
             prompt_args_template,
             user_prompt,
@@ -475,7 +475,7 @@ def build_interface(
 
 def main():
     args = parse_args()
-    omni_llm = None
+    omni = None
 
     model_name = "/".join(args.model.split("/")[-2:])
     assert model_name in SUPPORTED_MODELS, (
@@ -485,9 +485,9 @@ def main():
     # Register signal handlers for graceful shutdown
     def signal_handler(sig, frame):
         print("\nReceived interrupt signal, shutting down...")
-        if omni_llm is not None:
+        if omni is not None:
             try:
-                omni_llm.shutdown()
+                omni.shutdown()
             except Exception as e:
                 print(f"Error during shutdown: {e}")
         sys.exit(0)
@@ -495,21 +495,21 @@ def main():
     signal.signal(signal.SIGINT, signal_handler)
     signal.signal(signal.SIGTERM, signal_handler)
 
-    print(f"Initializing AsyncOmniLLM with model: {args.model}")
+    print(f"Initializing AsyncOmni with model: {args.model}")
     if args.stage_configs_path:
         print(f"Using custom stage configs: {args.stage_configs_path}")
 
     sampling_params = build_sampling_params(SEED, model_name)
-    omni_llm = AsyncOmniLLM(
+    omni = AsyncOmni(
         model=args.model,
         stage_configs_path=args.stage_configs_path,
         init_timeout=ASYNC_INIT_TIMEOUT,
     )
-    print("✓ AsyncOmniLLM initialized successfully")
+    print("✓ AsyncOmni initialized successfully")
     prompt_args_template = create_prompt_args(args)
 
     demo = build_interface(
-        omni_llm,
+        omni,
         sampling_params,
         prompt_args_template,
         args.model,
@@ -524,9 +524,9 @@ def main():
         print("\nShutting down...")
     finally:
         # Cleanup
-        if omni_llm is not None:
+        if omni is not None:
             try:
-                omni_llm.shutdown()
+                omni.shutdown()
             except Exception as e:
                 print(f"Error during cleanup: {e}")
 

examples/online_serving/qwen3_omni/README.md

@@ -45,7 +45,7 @@ This Web UI demo allows users to interact with the model through a web browser.
 
 ### Running Gradio Demo
 
-Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by
+Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmni by
 
 ```bash
 python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861

examples/online_serving/qwen3_omni/gradio_demo.py

@@ -15,7 +15,7 @@ from PIL import Image
 from vllm.assets.video import video_get_metadata, video_to_ndarrays
 from vllm.sampling_params import SamplingParams
 
-from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
+from vllm_omni.entrypoints.async_omni import AsyncOmni
 
 # Import utils from offline inference example
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../offline_inference/qwen3_omni"))
@@ -239,8 +239,8 @@ def process_video_file(
     return frames, metadata, audio_tuple
 
 
-async def run_inference_async_omni_llm(
-    omni_llm: AsyncOmniLLM,
+async def run_inference_async_omni(
+    omni: AsyncOmni,
     sampling_params: list[SamplingParams],
     prompt_args_template: SimpleNamespace,
     user_prompt: str,
@@ -249,7 +249,7 @@ async def run_inference_async_omni_llm(
     video_file: Optional[str] = None,
     use_audio_in_video: bool = False,
 ):
-    """Run inference using AsyncOmniLLM directly with multimodal support."""
+    """Run inference using AsyncOmni directly with multimodal support."""
     if not user_prompt.strip() and not audio_file and not image_file and not video_file:
         return "Please provide at least a text prompt or multimodal input.", None
 
@@ -331,7 +331,7 @@ async def run_inference_async_omni_llm(
         text_outputs: list[str] = []
         audio_output = None
 
-        async for stage_outputs in omni_llm.generate(
+        async for stage_outputs in omni.generate(
             prompt=omni_prompt,
             request_id=request_id,
             sampling_params_list=sampling_params,
@@ -366,12 +366,12 @@ async def run_inference_async_omni_llm(
 
 
 def build_interface(
-    omni_llm: AsyncOmniLLM,
+    omni: AsyncOmni,
     sampling_params: list[SamplingParams],
     prompt_args_template: SimpleNamespace,
     model: str,
 ):
-    """Build Gradio interface for AsyncOmniLLM mode."""
+    """Build Gradio interface for AsyncOmni mode."""
 
     async def run_inference(
         user_prompt: str,
@@ -380,8 +380,8 @@ def build_interface(
         video_file: Optional[str],
         use_audio_in_video: bool,
     ):
-        return await run_inference_async_omni_llm(
-            omni_llm,
+        return await run_inference_async_omni(
+            omni,
             sampling_params,
             prompt_args_template,
             user_prompt,
@@ -479,7 +479,7 @@ def build_interface(
 
 def main():
     args = parse_args()
-    omni_llm = None
+    omni = None
 
     model_name = "/".join(args.model.split("/")[-2:])
     assert model_name in SUPPORTED_MODELS, (
@@ -489,9 +489,9 @@ def main():
     # Register signal handlers for graceful shutdown
     def signal_handler(sig, frame):
         print("\nReceived interrupt signal, shutting down...")
-        if omni_llm is not None:
+        if omni is not None:
             try:
-                omni_llm.shutdown()
+                omni.shutdown()
             except Exception as e:
                 print(f"Error during shutdown: {e}")
         sys.exit(0)
@@ -499,21 +499,21 @@ def main():
     signal.signal(signal.SIGINT, signal_handler)
     signal.signal(signal.SIGTERM, signal_handler)
 
-    print(f"Initializing AsyncOmniLLM with model: {args.model}")
+    print(f"Initializing AsyncOmni with model: {args.model}")
     if args.stage_configs_path:
         print(f"Using custom stage configs: {args.stage_configs_path}")
 
     sampling_params = build_sampling_params(SEED, model_name)
-    omni_llm = AsyncOmniLLM(
+    omni = AsyncOmni(
         model=args.model,
         stage_configs_path=args.stage_configs_path,
         init_timeout=ASYNC_INIT_TIMEOUT,
     )
-    print("✓ AsyncOmniLLM initialized successfully")
+    print("✓ AsyncOmni initialized successfully")
     prompt_args_template = create_prompt_args(args)
 
     demo = build_interface(
-        omni_llm,
+        omni,
         sampling_params,
         prompt_args_template,
         args.model,
@@ -528,9 +528,9 @@ def main():
         print("\nShutting down...")
     finally:
         # Cleanup
-        if omni_llm is not None:
+        if omni is not None:
             try:
-                omni_llm.shutdown()
+                omni.shutdown()
             except Exception as e:
                 print(f"Error during cleanup: {e}")
 

mkdocs.yml

@@ -86,6 +86,7 @@ plugins:
       api_root_uri: "api"
       nav_item_prefix: ""  # No prefix in navigation tree (clean names)
       show_full_namespace: false  # Show only module name, not full path
+      on_implicit_namespace_package: skip  # Skip implicit namespace packages without warning
       exclude:
         - "re:vllm_omni\\._.*"  # Internal modules
         - "vllm_omni.diffusion.models.qwen_image"  # avoid importing vllm in mkdocs building

tests/test_omni_llm.py

@@ -342,14 +342,14 @@ def _setup_log_mocks(monkeypatch):
 def mock_get_config(monkeypatch):
     """Auto-mock get_config and related model loading functions to avoid model path validation."""
     # CRITICAL: Mock tokenizer-related imports FIRST, before any module imports
-    # This prevents ImportError when async_omni_llm is imported (which happens via omni_stage)
+    # This prevents ImportError when async_omni is imported (which happens via omni_stage)
     import sys
 
     fake_tokenizer = MagicMock()
     fake_tokenizer.encode = MagicMock(return_value=[1, 2, 3])
     fake_tokenizer.decode = MagicMock(return_value="test")
 
-    # Mock init_tokenizer_from_configs (used in async_omni_llm)
+    # Mock init_tokenizer_from_configs (used in async_omni)
     def _mock_init_tokenizer_from_configs(model_config=None, **kwargs):
         return fake_tokenizer
 
@@ -367,8 +367,9 @@ def mock_get_config(monkeypatch):
         tokenizer_module = sys.modules[tokenizer_module_path]
         setattr(tokenizer_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
 
-    # CRITICAL: Mock length_from_prompt_token_ids_or_embeds BEFORE trying to mock async_omni_llm
-    # This is because async_omni_llm imports processor.py, which imports this function at module level
+    # CRITICAL: Mock length_from_prompt_token_ids_or_embeds BEFORE trying to mock async_omni
+
+    # This is because async_omni imports processor.py, which imports this function at module level
     # Mock length_from_prompt_token_ids_or_embeds (used in processor.py)
     def _mock_length_from_prompt_token_ids_or_embeds(prompt_token_ids=None, prompt_embeds=None):
         # Return a reasonable default length
@@ -402,19 +403,19 @@ def mock_get_config(monkeypatch):
             processor_module, "length_from_prompt_token_ids_or_embeds", _mock_length_from_prompt_token_ids_or_embeds
         )
 
-    # Strategy 3: Now mock async_omni_llm AFTER length_from_prompt_token_ids_or_embeds is mocked
-    # This prevents ImportError when async_omni_llm imports processor.py
+    # Strategy 3: Now mock async_omni AFTER length_from_prompt_token_ids_or_embeds is mocked
+    # This prevents ImportError when async_omni imports processor.py
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.async_omni_llm.init_tokenizer_from_configs",
+        "vllm_omni.entrypoints.async_omni.init_tokenizer_from_configs",
         _mock_init_tokenizer_from_configs,
         raising=False,
     )
 
-    # Strategy 4: If async_omni_llm is already imported, patch it directly
-    async_omni_llm_path = "vllm_omni.entrypoints.async_omni_llm"
-    if async_omni_llm_path in sys.modules:
-        async_omni_llm_module = sys.modules[async_omni_llm_path]
-        setattr(async_omni_llm_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
+    # Strategy 4: If async_omni is already imported, patch it directly
+    async_omni_path = "vllm_omni.entrypoints.async_omni"
+    if async_omni_path in sys.modules:
+        async_omni_module = sys.modules[async_omni_path]
+        setattr(async_omni_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
 
     # Now mock get_config and other functions
     fake_hf_config = MagicMock()

vllm_omni/__init__.py

@@ -21,7 +21,7 @@ except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
     patch = None  # type: ignore
 
 from .config import OmniModelConfig
-from .entrypoints.async_omni_llm import AsyncOmniLLM
+from .entrypoints.async_omni import AsyncOmni
 
 # Main entry points
 from .entrypoints.omni import Omni
@@ -34,7 +34,7 @@ __all__ = [
     "__version_tuple__",
     # Main components
     "Omni",
-    "AsyncOmniLLM",
+    "AsyncOmni",
     # Configuration
     "OmniModelConfig",
     # All other components are available through their respective modules

vllm_omni/entrypoints/async_omni.py

@@ -51,8 +51,8 @@ from vllm_omni.outputs import OmniRequestOutput
 logger = init_logger(__name__)
 
 
-class AsyncOmniLLM(EngineClient):
-    """Async entry point for vLLM-Omni inference.
+class AsyncOmni(EngineClient):
+    """Async entry point for vLLM-omni inference.
 
     This class provides an asynchronous interface for running multi-modal
     comprehension and generation models. It orchestrates multiple
@@ -75,7 +75,7 @@ class AsyncOmniLLM(EngineClient):
         **kwargs: Additional keyword arguments passed to stage engines
 
     Example:
-        >>> async_llm = AsyncOmniLLM(model="Qwen/Qwen2.5-Omni-7B")
+        >>> async_llm = AsyncOmni(model="Qwen/Qwen2.5-Omni-7B")
         >>> async for output in async_llm.generate(
         ...     prompt="Hello",
         ...     request_id="req-1",
@@ -175,7 +175,7 @@ class AsyncOmniLLM(EngineClient):
 
         Sends shutdown signals to all stage input queues and stops
         all stage worker processes. This method should be called
-        when done using the AsyncOmniLLM instance.
+        when done using the AsyncOmni instance.
         """
         for q in self._stage_in_queues:
             try:
@@ -611,13 +611,13 @@ class AsyncOmniLLM(EngineClient):
         **kwargs,
     ):
         """Generate outputs for a request from a pooling model."""
-        raise NotImplementedError("encode() is not implemented for AsyncOmniLLM")
+        raise NotImplementedError("encode() is not implemented for AsyncOmni")
 
     async def start_profile(self) -> None:
-        raise NotImplementedError("start_profile() is not implemented for AsyncOmniLLM")
+        raise NotImplementedError("start_profile() is not implemented for AsyncOmni")
 
     async def stop_profile(self) -> None:
-        raise NotImplementedError("stop_profile() is not implemented for AsyncOmniLLM")
+        raise NotImplementedError("stop_profile() is not implemented for AsyncOmni")
 
 
 class AsyncOmniStageLLM(AsyncLLM):

vllm_omni/entrypoints/omni_stage.py

@@ -671,7 +671,7 @@ async def _stage_worker_async(
     import logging as _logging
     import time as _time
 
-    from vllm_omni.entrypoints.async_omni_llm import AsyncOmniStageLLM
+    from vllm_omni.entrypoints.async_omni import AsyncOmniStageLLM
     from vllm_omni.entrypoints.log_utils import (
         compute_and_log_stage_request_stats,
         count_tokens_from_outputs,

vllm_omni/entrypoints/openai/api_server.py

@@ -38,7 +38,7 @@ from vllm.logger import init_logger
 from vllm.transformers_utils.tokenizer import MistralTokenizer
 from vllm.utils import decorate_logs
 
-from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
+from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
 
 logger = init_logger(__name__)
@@ -65,7 +65,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
     if log_config is not None:
         uvicorn_kwargs["log_config"] = log_config
 
-    async with build_async_omni_llm(
+    async with build_async_omni(
         args,
         client_config=client_config,
     ) as engine_client:
@@ -108,15 +108,15 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
 
 
 @asynccontextmanager
-async def build_async_omni_llm(
+async def build_async_omni(
     args: Namespace,
     *,
     disable_frontend_multiprocessing: Optional[bool] = None,
     client_config: Optional[dict[str, Any]] = None,
 ) -> AsyncIterator[EngineClient]:
-    """Build an AsyncOmniLLM instance from command-line arguments.
+    """Build an AsyncOmni instance from command-line arguments.
 
-    Creates an async context manager that yields an AsyncOmniLLM instance
+    Creates an async context manager that yields an AsyncOmni instance
     configured from the provided arguments. Handles forkserver setup if
     needed and ensures proper cleanup on exit.
 
@@ -127,7 +127,7 @@ async def build_async_omni_llm(
         client_config: Optional client configuration dictionary
 
     Yields:
-        EngineClient instance (AsyncOmniLLM) ready for use
+        EngineClient instance (AsyncOmni) ready for use
     """
     if os.getenv("VLLM_WORKER_MULTIPROC_METHOD") == "forkserver":
         # The executor is expected to be mp.
@@ -138,25 +138,25 @@ async def build_async_omni_llm(
         forkserver.ensure_running()
         logger.debug("Forkserver setup complete!")
 
-    # Context manager to handle async_omni_llm lifecycle
+    # Context manager to handle async_omni lifecycle
     # Ensures everything is shutdown and cleaned up on error/exit
-    async with build_async_omni_llm_from_stage_config(
+    async with build_async_omni_from_stage_config(
         args,
         disable_frontend_multiprocessing=disable_frontend_multiprocessing,
-    ) as async_omni_llm:
-        yield async_omni_llm
+    ) as async_omni:
+        yield async_omni
 
 
 @asynccontextmanager
-async def build_async_omni_llm_from_stage_config(
+async def build_async_omni_from_stage_config(
     args: Namespace,
     *,
     disable_frontend_multiprocessing: bool = False,
     client_config: Optional[dict[str, Any]] = None,
 ) -> AsyncIterator[EngineClient]:
-    """Create AsyncOmniLLM from stage configuration.
+    """Create AsyncOmni from stage configuration.
 
-    Creates an AsyncOmniLLM instance either in-process or using multiprocess
+    Creates an AsyncOmni instance either in-process or using multiprocess
     RPC. Loads stage configurations from the model or from a specified path.
 
     Args:
@@ -166,7 +166,7 @@ async def build_async_omni_llm_from_stage_config(
         client_config: Optional client configuration dictionary
 
     Yields:
-        EngineClient instance (AsyncOmniLLM) ready for use
+        EngineClient instance (AsyncOmni) ready for use
 
     Note:
         Stage configurations are loaded from args.stage_configs_path if provided,
@@ -182,21 +182,21 @@ async def build_async_omni_llm_from_stage_config(
             "To disable frontend multiprocessing, set VLLM_USE_V1=0."
         )
 
-    async_omni_llm: Optional[EngineClient] = None
+    async_omni: Optional[EngineClient] = None
 
     try:
         if getattr(args, "stage_configs_path", None):
-            async_omni_llm = AsyncOmniLLM(model=args.model, stage_configs_path=args.stage_configs_path)
+            async_omni = AsyncOmni(model=args.model, stage_configs_path=args.stage_configs_path)
         else:
-            async_omni_llm = AsyncOmniLLM(model=args.model)
+            async_omni = AsyncOmni(model=args.model)
 
         # # Don't keep the dummy data in memory
         # await async_llm.reset_mm_cache()
 
-        yield async_omni_llm
+        yield async_omni
     finally:
-        if async_omni_llm:
-            async_omni_llm.shutdown()
+        if async_omni:
+            async_omni.shutdown()
 
 
 async def omni_init_app_state(
@@ -211,7 +211,7 @@ async def omni_init_app_state(
     and other server configuration needed for handling API requests.
 
     Args:
-        engine_client: Engine client instance (AsyncOmniLLM)
+        engine_client: Engine client instance (AsyncOmni)
         vllm_config: vLLM configuration object
         state: FastAPI application state object to initialize
         args: Parsed command-line arguments