[Entrypoint][Refactor] vLLM-Omni Entrypoint Refactoring (#1908)

Signed-off-by: yinpe <11810305@mail.sustech.edu.cn>
Signed-off-by: Peiqi Yin <yinpeiqi809@gmail.com>
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Signed-off-by: wuhang <whlbx@hotmail.com>
Signed-off-by: linyueqian <linyueqian@outlook.com>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: lishunyang <lishunyang12@163.com>
Signed-off-by: Chenguang ZHENG <645327136@qq.com>
Signed-off-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Co-authored-by: yinpe <11810305@mail.sustech.edu.cn>
Co-authored-by: Peiqi Yin <yinpeiqi809@gmail.com>
Co-authored-by: wuhang <whlbx@hotmail.com>
Co-authored-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Co-authored-by: linyueqian <linyueqian@outlook.com>
Co-authored-by: princepride <wangzhipeng628@gmail.com>
Co-authored-by: lishunyang <lishunyang12@163.com>

.github/ISSUE_TEMPLATE/400-bug-report.yml

@@ -74,28 +74,21 @@ body:
       If relevant, add a minimal example so that we can reproduce the error by running the code. It is very important for the snippet to be as succinct (minimal) as possible, so please take time to trim down any irrelevant code to help us debug efficiently. We are going to copy-paste your code and we expect to get the same result as you did: avoid any external data, and include the relevant imports, etc. For example:
 
       ```python
-      from vllm_omni import OmniLLM, create_ar_stage_config, create_dit_stage_config
+      from vllm_omni.entrypoints.omni import Omni
+      from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+      from vllm import SamplingParams
 
-      # Create stage configurations
-      ar_config = create_ar_stage_config(
-          stage_id=0,
-          model_path="Qwen/Qwen3-0.6B",
-          input_modalities=["text"],
-          output_modalities=["text"]
+      omni = Omni(
+          model="Qwen/Qwen-Image",
+          stage_configs_path="/path/to/stage_configs.yaml",
       )
 
-      dit_config = create_dit_stage_config(
-          stage_id=1,
-          model_path="stabilityai/stable-diffusion-2-1",
-          input_modalities=["text"],
-          output_modalities=["image"]
-      )
-
-      # Initialize OmniLLM
-      omni = OmniLLM([ar_config, dit_config])
-
-      # Generate
-      outputs = omni.generate(prompt="A scenic watercolor painting of a lighthouse at sunset")
+      prompts = [{"prompt": "A scenic watercolor painting of a lighthouse at sunset"}]
+      sampling_params_list = [
+          SamplingParams(max_tokens=1),
+          OmniDiffusionSamplingParams(num_outputs_per_prompt=1),
+      ]
+      outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list)
       ```
 
       If the code is too long (hopefully, it isn't), feel free to put it in a public gist and link it in the issue: https://gist.github.com.

benchmarks/qwen3-omni/README.md

@@ -56,9 +56,9 @@ What it does:
 - Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`.
 - Uses `benchmarks/build_dataset/top100.txt` and writes to:
   - Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
-    - `omni_llm_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
-    - `omni_llm_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
-    - `omni_llm_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
+    - `omni_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
+    - `omni_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
+    - `omni_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
   - Outputs: `benchmarks/qwen3-omni/vllm_omni/outputs/` — ~100 text and `.wav` files.
 
 Key checks:

benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh

@@ -26,12 +26,12 @@ else
                       --log-stats \
                       --log-dir $log_dir
     echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
-    echo "  - omni_llm_pipeline_text                       run dir/base name"
-    echo "  - omni_llm_pipeline_text.orchestrator.stats.jsonl  orchestrator-stage latency stats"
-    echo "  - omni_llm_pipeline_text.overall.stats.jsonl       overall latency/TPS stats"
-    echo "  - omni_llm_pipeline_text.stage0.log                per-stage detailed logs"
-    echo "  - omni_llm_pipeline_text.stage1.log"
-    echo "  - omni_llm_pipeline_text.stage2.log"
+    echo "  - omni_pipeline_text                       run dir/base name"
+    echo "  - omni_pipeline_text.orchestrator.stats.jsonl  orchestrator-stage latency stats"
+    echo "  - omni_pipeline_text.overall.stats.jsonl       overall latency/TPS stats"
+    echo "  - omni_pipeline_text.stage0.log                per-stage detailed logs"
+    echo "  - omni_pipeline_text.stage1.log"
+    echo "  - omni_pipeline_text.stage2.log"
     echo "Key checks: overall.stats.jsonl for end-to-end latency/TPS; orchestrator.stats.jsonl for stable per-stage latency; stage*.log for errors or long tails."
     echo "  - outputs/             Generated txt and wav files, there should be 100 text and wav files generated respectively"
 fi

docs/api/README.md

@@ -6,19 +6,13 @@ Main entry points for vLLM-Omni inference and serving.
 
 - [vllm_omni.entrypoints.async_omni.AsyncOmni][]
 - [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
-- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
 - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
 - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
 - [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
 - [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
 - [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
 - [vllm_omni.entrypoints.omni.Omni][]
-- [vllm_omni.entrypoints.omni.OmniBase][]
-- [vllm_omni.entrypoints.omni_diffusion.OmniDiffusion][]
-- [vllm_omni.entrypoints.omni_llm.OmniLLM][]
-- [vllm_omni.entrypoints.omni_stage.OmniStage][]
-- [vllm_omni.entrypoints.stage_utils.OmniStageTaskType][]
-- [vllm_omni.entrypoints.zmq_utils.ZmqQueue][]
+- [vllm_omni.entrypoints.omni_base.OmniBase][]
 
 ## Inputs
 
@@ -48,9 +42,7 @@ Engine classes for offline and online inference.
 - [vllm_omni.engine.OmniEngineCoreOutputs][]
 - [vllm_omni.engine.OmniEngineCoreRequest][]
 - [vllm_omni.engine.PromptEmbedsPayload][]
-- [vllm_omni.engine.arg_utils.AsyncOmniEngineArgs][]
 - [vllm_omni.engine.arg_utils.OmniEngineArgs][]
-- [vllm_omni.engine.input_processor.OmniInputProcessor][]
 - [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
 - [vllm_omni.engine.output_processor.OmniRequestState][]
 

docs/configuration/stage_configs.md

@@ -18,7 +18,7 @@ If users want to modify some part of it. The custom stage_configs file can be in
 For offline (Assume necessary dependencies have ben imported):
 ```python
 model_name = "Qwen/Qwen2.5-Omni-7B"
-omni_llm = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
+omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
 ```
 
 For online serving:
@@ -30,7 +30,7 @@ vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to
 
 Below is a specific example of stage_configs.yaml in Qwen2.5-omni.
 ```python
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
 stage_args:
   - stage_id: 0 # mark the unique id for each stage
     runtime: # The disaggregated configuration

docs/configuration/stage_configs/qwen2_5_omni.yaml

@@ -1,4 +1,4 @@
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
 stage_args:
   - stage_id: 0
     runtime:

docs/contributing/README.md

@@ -107,7 +107,7 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
 - `[CI/Build]` for build or continuous integration improvements.
 - `[Doc]` for documentation fixes and improvements.
 - `[Model]` for adding a new model or improving an existing model. Model name should appear in the title.
-- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `OmniLLM` class, etc.)
+- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `Omni`/`AsyncOmni`, etc.)
 - `[Kernel]` for changes affecting CUDA kernels or other compute kernels.
 - `[Core]` for changes in the core vLLM-Omni logic (e.g., `OmniProcessor`, `OmniARScheduler`, etc.)
 - `[Hardware][Vendor]` for hardware-specific changes. Vendor name should appear in the prefix, such as [Ascend] for Ascend NPUs.

docs/contributing/ci/CI_5levels.md

@@ -168,11 +168,6 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
-│   ├── omni.py                                 │   ├── test_omni.py                    ⬜  (E2E covered by e2e/offline, e2e/online)
-│   ├── omni_llm.py                             │   ├── test_omni_llm.py                ✅
-│   ├── omni_stage.py                            │   ├── test_omni_stage.py              ⬜  (partial in test_omni_stage_diffusion_config.py)
-│   ├── omni_diffusion.py                       │   ├── test_omni_diffusion.py          ✅
-│   ├── async_omni.py                            │   ├── test_async_omni.py              ✅ actually in e2e/online_serving/test_async_omni.py
 │   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)

docs/contributing/ci/tests_style.md

@@ -24,8 +24,6 @@ End-to-end tests verify the complete functionality of a system or component. For
 
 - **`tests/e2e/online_serving/`**: Tests for online serving scenarios (e.g., API server tests)
 
-**Example:** The test file for `vllm_omni/entrypoints/omni_llm.py` should be located at `tests/entrypoints/test_omni_llm.py`.
-
 ## Test Directory Structure
 
 The ideal directory structure mirrors the source code organization. Legend: `✅` = test exists, `⬜` = suggested to add.
@@ -75,11 +73,6 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
-│   ├── omni.py                                 │   ├── test_omni.py                    ⬜  (E2E covered by e2e/offline, e2e/online)
-│   ├── omni_llm.py                             │   ├── test_omni_llm.py                ✅
-│   ├── omni_stage.py                            │   ├── test_omni_stage.py              ⬜  (partial in test_omni_stage_diffusion_config.py)
-│   ├── omni_diffusion.py                       │   ├── test_omni_diffusion.py          ✅
-│   ├── async_omni.py                            │   ├── test_async_omni.py              ✅ actually in e2e/online_serving/test_async_omni.py
 │   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
@@ -170,7 +163,7 @@ vllm_omni/                                    tests/
 
 ### Naming Conventions
 
-- **Unit Tests**: Use `test_<module_name>.py` format. Example: `omni_llm.py` → `test_omni_llm.py`
+- **Unit Tests**: Use `test_<module_name>.py` format. Example: `stage_utils.py` → `test_stage_utils.py`
 
 - **E2E Tests**: Place in `tests/e2e/offline_inference/` or `tests/e2e/online_serving/` with descriptive names. Example: `tests/e2e/offline_inference/test_qwen3_omni.py`, `tests/e2e/offline_inference/test_diffusion_model.py`
 

docs/contributing/model/adding_omni_model.md

@@ -330,54 +330,46 @@ Stage transitions are the mechanism by which outputs from one stage are converte
 
 ### Where Stage Transitions Are Called
 
-Stage transitions happen automatically in the orchestrator (`OmniLLM` class) during the generation loop. Here's the detailed flow:
+Stage transitions happen automatically in the runtime orchestrator. Here's the detailed flow:
 
-1. **Location**: `vllm_omni/entrypoints/omni_llm.py` in the `_run_generation()` method
+1. **Location**: `vllm_omni/engine/orchestrator.py` in `_forward_to_next_stage()`
 2. **Trigger**: When a stage completes processing and produces outputs
 3. **Execution Flow**:
    ```python
-   # In omni_llm.py, _run_generation() method (around line 345-460)
+   # In orchestrator.py
+   next_stage_id = stage_id + 1
+   next_client = self.stage_clients[next_stage_id]
+   params = req_state.sampling_params_list[next_stage_id]
 
-   # Main orchestrator loop polls each stage for completed requests
-   for stage_id, stage in enumerate(self.stage_list):
-       result = stage.try_collect()  # Get completed request
-       if result is None:
-           continue
+   # Save current stage outputs so stage_input_processors can consume them.
+   self.stage_clients[stage_id].set_engine_outputs([output])
 
-       # Store outputs from this stage
-       engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
-       stage.set_engine_outputs(engine_outputs)
+   # THIS IS WHERE STAGE TRANSITION HAPPENS
+   next_inputs = next_client.process_engine_inputs(
+       stage_list=self.stage_clients,
+       prompt=req_state.prompt,
+   )
 
-       # Check if there's a next stage to forward to
-       next_stage_id = stage_id + 1
-       if next_stage_id < len(self.stage_list):
-           next_stage: OmniStage = self.stage_list[next_stage_id]
-
-           # THIS IS WHERE STAGE TRANSITION HAPPENS
-           next_inputs = next_stage.process_engine_inputs(
-               self.stage_list,
-               [request_id_to_prompt[req_id]]
-           )
-
-           # Submit to next stage
-           task = {
-               "type": OmniStageTaskType.GENERATE,
-               "request_id": req_id,
-               "engine_inputs": next_inputs[0],
-               "sampling_params": sampling_params_list[next_stage_id],
-           }
-           next_stage.submit(task)
+   # Build and submit request(s) to the next stage.
+   for next_input in next_inputs:
+       request = build_engine_core_request_from_tokens(
+           request_id=req_id,
+           prompt=next_input,
+           params=params,
+           model_config=self.stage_vllm_configs[next_stage_id].model_config,
+       )
+       await next_client.add_request_async(request)
    ```
 
 ### How Stage Transitions Work
 
 The stage transition process follows these steps:
 
-1. **Stage Completion**: When a stage finishes processing a request, it stores outputs via `stage.set_engine_outputs(engine_outputs)`
+1. **Stage Completion**: When a stage finishes processing a request, the orchestrator stores outputs via `stage_client.set_engine_outputs(...)`
 
 2. **Transition Detection**: The orchestrator checks if there's a next stage and calls `process_engine_inputs()` on it
 
-3. **Input Processing**: The `process_engine_inputs()` method in `OmniStage` (`omni_stage.py`) handles the transition:
+3. **Input Processing**: The stage input processor configured in stage YAML (under `vllm_omni/model_executor/stage_input_processors/`) handles the transition:
    ```python
    def process_engine_inputs(
        self, stage_list: list[Any], prompt: OmniTokensPrompt | TextPrompt = None

docs/contributing/profiling.md

@@ -23,30 +23,32 @@ export VLLM_PROFILER_MAX_ITERS=1
 The profiler is default to function across all stages. But It is highly recommended to profile specific stages by passing the stages list, preventing from producing too large trace files:
 ```python
 # Profile all stages
-omni_llm.start_profile()
+omni.start_profile()
 
 # Only profile Stage 1
-omni_llm.start_profile(stages=[1])
+omni.start_profile(stages=[1])
 ```
 
 ```python
 # Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
-omni_llm.start_profile(stages=[0, 2])
+omni.start_profile(stages=[0, 2])
 ```
 
 **Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
 
 ```python
-from vllm_omni import omni_llm
+from vllm_omni.entrypoints.omni import Omni
+
+omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
 
 profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
 
 # 1. Start profiling if enabled
 if profiler_enabled:
-    omni_llm.start_profile(stages=[0])
+    omni.start_profile(stages=[0])
 
 # Initialize generator
-omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
 
 total_requests = len(prompts)
 processed_count = 0
@@ -57,21 +59,21 @@ for stage_outputs in omni_generator:
     # ... [Output processing logic for text/audio would go here] ...
 
     # Update count to track when to stop profiling
-    processed_count += len(stage_outputs.request_output)
+    processed_count += 1
 
     # 2. Check if all requests are done to stop the profiler safely
     if profiler_enabled and processed_count >= total_requests:
         print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
 
         # Stop the profiler while workers are still active
-        omni_llm.stop_profile()
+        omni.stop_profile()
 
         # Wait for traces to flush to disk
         print("[Info] Waiting 30s for workers to write trace files to disk...")
         time.sleep(30)
         print("[Info] Trace export wait time finished.")
 
-omni_llm.close()
+omni.close()
 ```
 
 

docs/design/architecture_overview.md

@@ -67,12 +67,12 @@ According to analysis for current popular open-source models, most of them have
 | Component         | Description                                                                                                                              |
 | ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
 | **OmniRouter**    | provide an intelligent router for Omni-modality requests dispatch                                                                        |
-| **EntryPoints**   | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni) and provide the OmniStage abstraction for different AR/DiT stages |
+| **EntryPoints**   | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni), while `AsyncOmniEngine` and `Orchestrator` coordinate multi-stage AR/DiT execution |
 | **AR**            | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management                                 |
 | **Diffusion**     | natively implemented and optimized using acceleration components                                                                         |
 | **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages                    |
 
-Disaggregated stages are managed through configuration, such as in the Qwen3-Omni example, where stages like Thinker, Talker, and Code2wav are defined as separate OmniStage instances with specific resources and input/output type.
+Disaggregated stages are managed through stage configuration. In Qwen3-Omni, Thinker/Talker/Code2wav are declared as separate configured stages, and runtime routing is handled by `Orchestrator` over `StageEngineCoreClient` / `StageDiffusionClient`.
 
 ## Main features
 
@@ -127,10 +127,10 @@ Taking **Qwen3-Omni** as an example:
 The **Omni** class provides a Python interface for offline batched inference. Users initialize the Omni class with a Hugging Face model name and use the generate method, passing inputs that include both text prompts and multi-modal data:
 
 ```
-# Create an omni_lm with HF model name.
+# Create an omni runtime with HF model name.
 from vllm_omni.entrypoints.omni import Omni
 
-omni_lm = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
+omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
 
 # Example prompts.
 om_inputs = {"prompt": prompt,
@@ -140,7 +140,7 @@ om_inputs = {"prompt": prompt,
              }}
 
 # Generate texts and audio from the multi-modality inputs.
-outputs = omni_lm.generate(om_inputs, sampling_params_list)
+outputs = omni.generate(om_inputs, sampling_params_list)
 ```
 
 ## Online Serving

docs/design/module/ar_module.md

@@ -85,9 +85,6 @@ classDiagram
     class InputProcessor {
         +process_inputs() EngineCoreRequest
     }
-    class OmniInputProcessor {
-        +process_inputs() OmniEngineCoreRequest
-    }
 
     class VLLMOutputProcessor {
         +process_outputs() OutputProcessorOutput
@@ -96,7 +93,6 @@ classDiagram
         +process_outputs() OutputProcessorOutput
         +_route_and_normalize()
     }
-    InputProcessor <|-- OmniInputProcessor
     VLLMOutputProcessor <|-- MultimodalOutputProcessor
 ```
 
@@ -105,7 +101,7 @@ classDiagram
 - **Scheduler**: `OmniARScheduler` extends `vllm.v1.core.sched.scheduler.Scheduler` to enrich scheduled requests with omni-specific payloads
 - **Worker**: `GPUARWorker` extends `vllm.v1.worker.gpu_worker.Worker` to initialize AR-specific model runners
 - **ModelRunner**: `GPUARModelRunner` extends `OmniGPUModelRunner` → `vllm.v1.worker.gpu_model_runner.GPUModelRunner` to expose hidden states and handle multimodal outputs
-- **InputProcessor**: `OmniInputProcessor` extends `vllm.v1.engine.input_processor.InputProcessor` to serialize prompt embeddings and additional information
+- **InputProcessor**: Stage-0 uses upstream `vllm.v1.engine.input_processor.InputProcessor`; `AsyncOmniEngine` then restores omni-specific payloads (for example `additional_information` and `prompt_embeds`) when building `OmniEngineCoreRequest`
 - **OutputProcessor**: `MultimodalOutputProcessor` extends `vllm.v1.engine.output_processor.OutputProcessor` to route and accumulate multimodal outputs
 
 ## 3. Scheduler Design
@@ -118,7 +114,7 @@ The following diagram illustrates the request flow through the AR module compone
 
 ```mermaid
 flowchart TD
-    A[OmniInputProcessor] -->|OmniEngineCoreRequest| B[OmniARScheduler]
+    A[InputProcessor stage-0 in AsyncOmniEngine] -->|EngineCoreRequest then upgraded to OmniEngineCoreRequest| B[OmniARScheduler]
     B -->|schedule: OmniNewRequestData| C[GPUARWorker]
     C -->|SchedulerOutput| D[GPUARModelRunner]
     D -->|execute_model: None| E[Model Forward Pass]
@@ -297,15 +293,17 @@ The input/output processing pipeline handles serialization, routing, and accumul
 ```mermaid
 sequenceDiagram
     participant Client
-    participant OmniInputProcessor
+    participant AsyncOmniEngine
+    participant InputProcessor
     participant Scheduler
     participant ModelRunner
     participant MultimodalOutputProcessor
     participant Client
 
-    Client->>OmniInputProcessor: prompt + prompt_embeds + additional_info
-    OmniInputProcessor->>OmniInputProcessor: Serialize tensors to payloads
-    OmniInputProcessor->>Scheduler: OmniEngineCoreRequest (with payloads)
+    Client->>AsyncOmniEngine: prompt + prompt_embeds + additional_info
+    AsyncOmniEngine->>InputProcessor: process_inputs()
+    InputProcessor->>Scheduler: EngineCoreRequest
+    AsyncOmniEngine->>AsyncOmniEngine: _upgrade_to_omni_request() + serialize_additional_information()
     Scheduler->>ModelRunner: OmniNewRequestData (with payloads)
     ModelRunner->>ModelRunner: Decode payloads → CPU tensors
     ModelRunner->>ModelRunner: Overlay prompt_embeds on inputs_embeds
@@ -318,43 +316,18 @@ sequenceDiagram
     MultimodalOutputProcessor->>Client: RequestOutput (with multimodal_output)
 ```
 
-### OmniInputProcessor
+### Stage-0 Input Processing
 
-`OmniInputProcessor` extends the base input processor to serialize prompt embeddings and additional information for inter-stage transfer.
-
-#### Payload Serialization
-
-Converts PyTorch tensors to serialized payloads:
+Stage-0 now uses upstream `InputProcessor` directly, and `AsyncOmniEngine` upgrades the request to `OmniEngineCoreRequest` while restoring omni-specific payloads.
 
 ```python
-def process_inputs(self, ...) -> OmniEngineCoreRequest:
-    # Serialize prompt_embeds
-    if "prompt_embeds" in decoder_inputs:
-        pe_cpu = decoder_inputs["prompt_embeds"].detach().to("cpu").contiguous()
-        prompt_embeds_payload = PromptEmbedsPayload(
-            data=pe_cpu.numpy().tobytes(),
-            shape=[seq_len, hidden_size],
-            dtype=dtype_str,
-        )
-
-    # Serialize additional_information
-    if "additional_information" in decoder_inputs:
-        entries = {}
-        for key, value in raw_info.items():
-            if isinstance(value, torch.Tensor):
-                entry = AdditionalInformationEntry(
-                    tensor_data=value.numpy().tobytes(),
-                    tensor_shape=list(value.shape),
-                    tensor_dtype=dtype_str,
-                )
-            entries[key] = entry
-        additional_information_payload = AdditionalInformationPayload(entries=entries)
-
-    return OmniEngineCoreRequest(
-        # ... standard fields ...
-        prompt_embeds=prompt_embeds_payload,
-        additional_information=additional_information_payload,
-    )
+request = self.input_processor.process_inputs(
+    request_id=request_id,
+    prompt=prompt,
+    params=params,
+    supported_tasks=self.supported_tasks,
+)
+request = _upgrade_to_omni_request(request, prompt)
 ```
 
 ### MultimodalOutputProcessor
@@ -400,13 +373,13 @@ The AR module of vLLM-Omni extends vLLM through strategic inheritance and minima
 ### Key Design Patterns
 
 1. **Inheritance over composition**: Extends vLLM classes to preserve compatibility with existing scheduling, batching, and execution mechanisms
-2. **Payload serialization**: Uses `PromptEmbedsPayload` and `AdditionalInformationPayload` for efficient inter-stage data transfer
+2. **Payload serialization**: Uses serialized `additional_information` payloads together with prompt-embedding handoff for efficient inter-stage transfer
 3. **Two-phase execution**: Maintains vLLM's execute/sample separation for AR models while supporting single-phase execution for generation models
 4. **Multimodal routing**: Routes outputs by `output_type` and accumulates tensors incrementally to support streaming
 
 ### Differences from vLLM
 
-- **Payload support**: Serialized prompt embeddings and additional information enable direct transfer between pipeline stages
+- **Payload support**: Serialized additional information and prompt embeddings enable direct transfer between pipeline stages
 - **Multimodal handling**: Extended input/output processors support images, audio, and other modalities alongside text
 - **Hidden state exposure**: AR model runners expose per-request hidden states via `pooler_output` for downstream consumption
 - **Generation scheduler**: Fast-path scheduling for basic heterogeneous architectures that complete in one step

docs/design/module/async_omni_architecture.md

@@ -0,0 +1,203 @@
+# AsyncOmni Architecture (Qwen3-Omni Example)
+
+## 1. System Architecture
+
+```text
+• ┌─────────────────────────────────────────────────────────────────────────────────┐
+  │                                    API Layer                                    │
+  │  ┌─────────────────────────────────────┐  ┌──────────────────────────────────┐  │
+  │  │ AsyncOmni (EngineClient)            │  │ Omni                             │  │
+  │  │ • generate() / abort() / shutdown() │  │ • generate()                     │  │
+  │  │ • _final_output_handler()           │  │                                  |  │
+  │  └─────────────────────────────────────┘  └──────────────────────────────────┘  │
+  ├─────────────────────────────────────────────────────────────────────────────────┤
+  │                              Engine Layer (Proxy)                               │
+  │  ┌───────────────────────────────────────────────────────────────────────────┐  │
+  │  │ AsyncOmniEngine                                                           │  │
+  │  │ • _bootstrap_orchestrator() & _initialize_stages()                        │  │
+  │  │ • add_request() / add_request_async() -> input_processor.process_inputs() │  │
+  │  │ • try_get_output() / try_get_output_async()                               │  │
+  │  └───────────────────┬─────────────────────────────────▲─────────────────────┘  │
+  │         request_queue (janus.Queue)        output_queue (janus.Queue)           │
+  ├──────────────────────┼─────────────────────────────────┼────────────────────────┤
+  │                      ▼        Orchestration Layer      │                        │
+  │  ┌───────────────────────────────────────────────────────────────────────────┐  │
+  │  │ Orchestrator [background thread]                                          │  │
+  │  │ • _request_handler()                                                      │  │
+  │  │     -  stage_client.add_request_async() & _prewarm_async_chunk_stages()   │  │
+  │  │ • _orchestration_output_handler()                                         │  │
+  │  │     -  _process_stage_outputs() -> output_processors[i].process_outputs() │  │
+  │  │     -  _route_output() & _forward_to_next_stage()                         │  │
+  │  └──────────┬─────────────────────────┬────────────────────────┬─────────────┘  │
+  ├─────────────┼─────────────────────────┼────────────────────────┼────────────────┤
+  │             │                 Communication Layer              │                │
+  │  ┌───────────────────────┐ ┌───────────────────────┐ ┌───────────────────────┐  │
+  │  │ StageEngineCoreClient │ │ StageEngineCoreClient │ │ StageDiffusionClient  │  │
+  │  │ • ZMQ ROUTER / PULL   │ │ • ZMQ ROUTER / PULL   │ │ • ZMQ ROUTER / PULL   │  │
+  │  │ • Msgpack codec       │ │ • Msgpack codec       │ │ • Msgpack codec       │  │
+  │  └──────────┬────────────┘ └──────────┬────────────┘ └──────────┬────────────┘  │
+  │             ▼ ZMQ IPC                 ▼ ZMQ IPC                 ▼ ZMQ IPC       │
+  ├─────────────────────────────────────────────────────────────────────────────────┤
+  │                                 Execution Layer                                 │
+  │  ┌───────────────────────┐ ┌───────────────────────┐ ┌───────────────────────┐  │
+  │  │ StageCoreProc         │ │ StageCoreProc         │ │ DiffusionEngine       │  │
+  │  │ [background process]  │ │ [background process]  │ │ [background process]  │  │
+  │  └───────────────────────┘ └───────────────────────┘ └───────────────────────┘  │
+  └─────────────────────────────────────────────────────────────────────────────────┘
+```
+
+## 2. Execution Flow (Arrow Steps, one generate request)
+
+```text
+[1] App
+    -> AsyncOmni.generate(prompt, request_id)
+
+[2] AsyncOmni
+    -> _final_output_handler()   (started on first request)
+    -> AsyncOmniEngine.add_request(stage_id=0, ...)
+
+[3] AsyncOmniEngine.add_request
+    -> (if stage-0 is llm and input is not EngineCoreRequest)
+       InputProcessor.process_inputs()
+       OutputProcessor[0].add_request()
+    -> request_queue.put(add_request_msg)
+
+[4] Orchestrator._request_handler
+    -> _handle_add_request(msg)
+    -> stage_clients[0].add_request_async(...)
+
+[5] Orchestrator._orchestration_loop (loop)
+    -> poll stage output
+       - llm stage: await get_output_async()
+       - diffusion stage: get_diffusion_output_async()
+    -> (llm stage) output_processors[i].process_outputs(...)
+    -> _route_output(...)
+    -> if finished and not final_stage and non-async-chunk:
+         _forward_to_next_stage(...)
+         -> next_stage.add_request_async(...)
+    -> output_queue.put(output)
+
+[6] AsyncOmni._final_output_loop (background coroutine)
+    -> AsyncOmniEngine.try_get_output_async()
+    -> route by request_id to ClientRequestState.queue
+
+[7] AsyncOmni._process_orchestrator_results
+    -> read from ClientRequestState.queue
+    -> _process_single_result(...)
+    -> yield OmniRequestOutput
+
+[8] Exit condition
+    -> receive result["finished"] == True
+    -> generate() ends
+```
+
+## 3. Runtime Sequence (one generate request)
+
+```mermaid
+sequenceDiagram
+    participant APP as App
+    participant AO as AsyncOmni
+    participant ENG as AsyncOmniEngine
+    participant ORCH as Orchestrator
+    participant S0 as Stage-0 Client
+    participant SN as Next Stage Client
+
+    APP->>AO: generate
+    AO->>AO: start output_handler once
+    AO->>ENG: add_request(stage_id=0, ...)
+    ENG->>ENG: input_processor.process_inputs()
+    ENG->>ORCH: request_queue.put(add_request)
+
+    ORCH->>ORCH: _handle_add_request
+    ORCH->>S0: add_request_async
+
+    loop poll route forward
+        ORCH->>S0: get_output_async / get_diffusion_output_async
+        ORCH->>ORCH: _route_output
+        alt need forward to next stage
+            ORCH->>SN: add_request_async
+        end
+        ORCH-->>ENG: output_queue.put
+    end
+
+    AO->>ENG: try_get_output_async
+    ENG-->>AO: message
+    AO-->>APP: yield OmniRequestOutput
+```
+
+## 4. Comparison
+
+Previous topology (reference):
+
+```text
+┌────────────────────────────────────────────────────────────────────────────┐
+│ Main Process                                                               │
+│  ┌──────────────────────┐   ┌────────────────────────────────────────────┐ │
+│  │ generate()           │   │ final_output_handler()                     │ │
+│  └──────────────────────┘   └────────────────────────────────────────────┘ │
+└──────────┬─────────────────────────┬─────────────────────────┬─────────────┘
+  mp.Queue (in_q/out_q)    mp.Queue (in_q/out_q)    mp.Queue (in_q/out_q)
+           ▼▲                        ▼▲                        ▼▲
+┌───────────────────────┐  ┌───────────────────────┐  ┌──────────────────────┐
+│ Worker Proc-0         │  │ Worker Proc-1         │  │ Worker Proc-2        │
+│ (Thinker LLM)         │  │ (Talker LLM)          │  │ (Vocoder)            │
+│  ┌────────────────┐   │  │  ┌────────────────┐   │  │  ┌────────────────┐  │
+│  │_stage_worker   │   │  │  │_stage_worker   │   │  │  │_stage_worker   │  │
+│  │_async()        │   │  │  │_async()        │   │  │  │_async()        │  │
+│  └────────────────┘   │  │  └────────────────┘   │  │  └────────────────┘  │
+│  ┌────────────────┐   │  │  ┌────────────────┐   │  │  ┌────────────────┐  │
+│  │output_handler()│   │  │  │output_handler()│   │  │  │output_handler()│  │
+│  └────────────────┘   │  │  └────────────────┘   │  │  └────────────────┘  │
+└──────────┬────────────┘  └──────────┬────────────┘  └──────────┬───────────┘
+       ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ
+┌──────────────────────┐   ┌──────────────────────┐   ┌──────────────────────┐
+│ EngineCore Proc-0    │   │ EngineCore Proc-1    │   │ EngineCore Proc-2    │
+│ (Thinker)            │   │ (Talker)             │   │ (Vocoder)            │
+└──────────────────────┘   └──────────────────────┘   └──────────────────────┘
+```
+
+Current topology:
+
+```text
+┌────────────────────────────────────────────────────────────────────────────┐
+│ Main Process                                                               │
+│  ┌──────────────────────────────────────────────────────────────────────┐  │
+│  │ Main Thread                                                          │  │
+│  │  ┌──────────────────────┐   ┌─────────────────────────────────────┐  │  │
+│  │  │ generate()           │   │ final_output_handler()              │  │  │
+│  │  └──────────────────────┘   └─────────────────────────────────────┘  │  │
+│  └──────────────────────────────────────────────────────────────────────┘  │
+│         janus.Queue (request_queue) ▼  ▲ janus.Queue (output_queue)        │
+│  ┌──────────────────────────────────────────────────────────────────────┐  │
+│  │ Orchestrator Thread                                                  │  │
+│  │  ┌──────────────────────┐  ┌──────────────────────────────────────┐  │  │
+│  │  │ _request_handler()   │  │ _orchestration_output_handler()      │  │  │
+│  │  └──────────────────────┘  └──────────────────────────────────────┘  │  │
+│  │  ┌────────────────────────────────────────────────────────────────┐  │  │
+│  │  │ _orchestration_loop(): poll/process/route outputs for all stages│  │  │
+│  │  └────────────────────────────────────────────────────────────────┘  │  │
+│  └───────┬─────────────────────────┬─────────────────────────┬──────────┘  │
+└──────────┬─────────────────────────┬─────────────────────────┬─────────────┘
+       ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ  
+  ┌──────────────────────┐  ┌──────────────────────┐  ┌──────────────────────┐
+  │ EngineCore Proc-0    │  │ EngineCore Proc-1    │  │ EngineCore Proc-2    │
+  │ (Thinker)            │  │ (Talker)             │  │ (Vocoder)            │
+  └──────────────────────┘  └──────────────────────┘  └──────────────────────┘
+```
+
+
+Test scripts:
+```bash
+# enter offline inference folder.
+cd examples/offline_inference/qwen2_5_omni
+python end2end.py --output-dir output_audio --query-type use_mixed_modalities
+
+cd ../qwen3_omni
+python end2end.py --output-dir output_audio --query-type text --async-chunk --enable-stats
+
+cd ../bagel
+python end2end.py --prompts "A cute cat"
+
+cd ../text_to_image
+python text_to_image.py --prompt "a cup of coffee on the table" --output output.png
+```

docs/features/custom_pipeline.md

@@ -97,7 +97,7 @@ outputs = omni.generate(
 )
 
 # Access custom trajectory data
-output = outputs[0].request_output[0]
+output = outputs[0].request_output
 print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
 print(f"Trajectory latents shape: {output.latents.shape}")
 ```

docs/getting_started/quickstart.md

@@ -42,7 +42,7 @@ if __name__ == "__main__":
     omni = Omni(model="Tongyi-MAI/Z-Image-Turbo")
     prompt = "a cup of coffee on the table"
     outputs = omni.generate(prompt)
-    images = outputs[0].request_output[0].images
+    images = outputs[0].request_output.images
     images[0].save("coffee.png")
 ```
 
@@ -70,7 +70,7 @@ if __name__ == "__main__":
     ]
     omni_outputs = omni.generate(prompts)
     for i_prompt, prompt_output in enumerate(omni_outputs):
-        this_request_output = prompt_output.request_output[0]
+        this_request_output = prompt_output.request_output
         this_images = this_request_output.images
         for i_image, image in enumerate(this_images):
             image.save(f"p{i_prompt}-img{i_image}.jpg")

docs/user_guide/diffusion/quantization/gguf.md

@@ -22,10 +22,16 @@
 CLI (examples/offline_inference/text_to_image/text_to_image.py)
   |
   v
-Omni (vllm_omni/entrypoints/omni.py)
+Omni (vllm_omni/entrypoints/__init__.py)
   |
   v
-OmniStage (diffusion)
+AsyncOmniEngine
+  |
+  v
+Orchestrator
+  |
+  v
+StageDiffusionClient
   |
   v
 DiffusionWorker
@@ -64,7 +70,13 @@ _generate_with_async_omni
 AsyncOmni
   |
   v
-DiffusionEngine
+AsyncOmniEngine
+  |
+  v
+Orchestrator
+  |
+  v
+StageDiffusionClient
   |
   v
 OmniRequestOutput

docs/user_guide/examples/offline_inference/text_to_image.md

@@ -19,7 +19,7 @@ if __name__ == "__main__":
     omni = Omni(model="Qwen/Qwen-Image")
     prompt = "a cup of coffee on the table"
     outputs = omni.generate(prompt)
-    images = outputs[0].request_output[0].images
+    images = outputs[0].request_output.images
     images[0].save("coffee.png")
 ```
 
@@ -37,7 +37,7 @@ if __name__ == "__main__":
     ]
     outputs = omni.generate(prompts)
     for i, output in enumerate(outputs):
-      image = output.request_output[0].images[0].save(f"{i}.jpg")
+      image = output.request_output.images[0].save(f"{i}.jpg")
 ```
 
 !!! info
@@ -72,7 +72,7 @@ if __name__ == "__main__":
       }
     ])
     for i, output in enumerate(outputs):
-      image = output.request_output[0].images[0].save(f"{i}.jpg")
+      image = output.request_output.images[0].save(f"{i}.jpg")
 ```
 
 ## Local CLI Usage

examples/offline_inference/bagel/end2end.py

@@ -34,7 +34,7 @@ def parse_args():
         help="Path to input image for img2img.",
     )
 
-    # OmniLLM init args
+    # Omni runtime init args
     parser.add_argument("--log-stats", action="store_true", default=False)
     parser.add_argument("--init-sleep-seconds", type=int, default=20)
     parser.add_argument("--batch-timeout", type=int, default=5)
@@ -163,25 +163,17 @@ def main():
 
     omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
 
-    for i, req_output in enumerate(omni_outputs):
+    img_idx = 0
+    for req_output in omni_outputs:
         images = getattr(req_output, "images", None)
-        if not images and hasattr(req_output, "output"):
-            if isinstance(req_output.output, list):
-                images = req_output.output
-            else:
-                images = [req_output.output]
+        if not images:
+            continue
 
-        if images:
-            for j, img in enumerate(images):
-                img.save(f"output_{i}_{j}.png")
-
-        if hasattr(req_output, "request_output") and req_output.request_output:
-            for stage_out in req_output.request_output:
-                if hasattr(stage_out, "images") and stage_out.images:
-                    for k, img in enumerate(stage_out.images):
-                        save_path = f"output_{i}_stage_{getattr(stage_out, 'stage_id', '?')}_{k}.png"
-                        img.save(save_path)
-                        print(f"[Info] Saved stage output image to {save_path}")
+        for j, img in enumerate(images):
+            save_path = f"output_{img_idx}_{j}.png"
+            img.save(save_path)
+            print(f"[Info] Saved image to {save_path}")
+        img_idx += 1
 
     print(omni_outputs)
 

examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py

@@ -167,27 +167,26 @@ def run_e2e():
     print(f"Received {len(outputs)} outputs.")
     for i, output in enumerate(outputs):
         try:
-            ro_list = output.request_output or []
-            if not ro_list:
+            ro = output.request_output
+            if ro is None:
                 print("No request_output found.")
                 continue
 
-            for ro in ro_list:
-                # Multimodal output may be attached to RequestOutput or CompletionOutput.
-                mm = getattr(ro, "multimodal_output", None)
-                if not mm and ro.outputs:
-                    mm = getattr(ro.outputs[0], "multimodal_output", None)
+            # Multimodal output may be attached to RequestOutput or CompletionOutput.
+            mm = getattr(ro, "multimodal_output", None)
+            if not mm and ro.outputs:
+                mm = getattr(ro.outputs[0], "multimodal_output", None)
 
-                if mm:
-                    print(f"Multimodal output keys: {mm.keys()}")
-                    if "audio" in mm:
-                        audio_out = mm["audio"]
-                        print(f"Generated Audio Shape: {audio_out.shape}")
-                        out_path = f"output_{i}.wav"
-                        sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050)
-                        print(f"Saved audio to {out_path}")
-                else:
-                    print("No multimodal output found.")
+            if mm:
+                print(f"Multimodal output keys: {mm.keys()}")
+                if "audio" in mm:
+                    audio_out = mm["audio"]
+                    print(f"Generated Audio Shape: {audio_out.shape}")
+                    out_path = f"output_{i}.wav"
+                    sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050)
+                    print(f"Saved audio to {out_path}")
+            else:
+                print("No multimodal output found.")
         except Exception as e:
             print(f"Error inspecting output: {e}")
     omni.close()

examples/offline_inference/custom_pipeline/image_to_image/image_edit.py

@@ -209,7 +209,7 @@ async def main():
     if not outputs:
         raise ValueError("No output produced from omni.generate()")
 
-    first_out = outputs[0].request_output[0]
+    first_out = outputs[0].request_output
     req_out: OmniRequestOutput = first_out
 
     # Verify trajectory data (from custom pipeline)

examples/offline_inference/fish_speech/end2end.py

@@ -129,8 +129,14 @@ def main(args):
 
     t_start = time.perf_counter()
     for stage_outputs in omni.generate(inputs):
-        for output in stage_outputs.request_output:
-            _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
+        request_output = stage_outputs.request_output
+        if request_output is None or not request_output.outputs:
+            continue
+        _save_wav(
+            output_dir,
+            request_output.request_id,
+            request_output.outputs[0].multimodal_output,
+        )
     t_end = time.perf_counter()
     logger.info("Total inference time: %.1f ms", (t_end - t_start) * 1000)
 

examples/offline_inference/glm_image/end2end.py

@@ -310,36 +310,35 @@ def main(args: argparse.Namespace) -> None:
 
     output_count = 0
     for stage_outputs in omni.generate(prompts, sampling_params_list, py_generator=True):
+        output = stage_outputs.request_output
         if stage_outputs.final_output_type == "image":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
+            request_id = output.request_id
 
-                # Get generated images
-                images = output.images if hasattr(output, "images") else []
-                if not images and hasattr(output, "multimodal_output"):
-                    images = output.multimodal_output.get("images", [])
+            # Get generated images
+            images = output.images if hasattr(output, "images") else []
+            if not images and hasattr(output, "multimodal_output"):
+                images = output.multimodal_output.get("images", [])
 
-                # Save each generated image
-                for idx, img in enumerate(images):
-                    if args.num_prompts == 1 and len(images) == 1:
-                        output_path = args.output
-                    else:
-                        base, ext = os.path.splitext(args.output)
-                        output_path = f"{base}_{request_id}_{idx}{ext}"
+            # Save each generated image
+            for idx, img in enumerate(images):
+                if args.num_prompts == 1 and len(images) == 1:
+                    output_path = args.output
+                else:
+                    base, ext = os.path.splitext(args.output)
+                    output_path = f"{base}_{request_id}_{idx}{ext}"
 
-                    if isinstance(img, Image.Image):
-                        save_image(img, output_path)
-                    else:
-                        print(f"Warning: Unexpected image type for request {request_id}: {type(img)}")
+                if isinstance(img, Image.Image):
+                    save_image(img, output_path)
+                else:
+                    print(f"Warning: Unexpected image type for request {request_id}: {type(img)}")
 
-                    output_count += 1
+                output_count += 1
 
         elif stage_outputs.final_output_type == "text":
             # AR stage output (intermediate, for debugging)
             if args.verbose:
-                for output in stage_outputs.request_output:
-                    print(f"AR output for request {output.request_id}:")
-                    print(f"  Token count: {len(output.outputs[0].token_ids)}")
+                print(f"AR output for request {output.request_id}:")
+                print(f"  Token count: {len(output.outputs[0].token_ids)}")
 
     gen_time = time.time() - gen_start_time
     print(f"\nGeneration completed in {gen_time:.2f}s")

examples/offline_inference/helios/end2end.py

@@ -293,12 +293,11 @@ def main():
                 )
 
             if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
-                if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
-                    inner_output = first_item.request_output[0]
-                    if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
-                        frames = inner_output.images[0] if inner_output.images else None
-                        if frames is None:
-                            raise ValueError("No video frames found in output.")
+                inner_output = first_item.request_output
+                if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
+                    frames = inner_output.images[0] if inner_output.images else None
+                    if frames is None:
+                        raise ValueError("No video frames found in output.")
             elif hasattr(first_item, "images") and first_item.images:
                 frames = first_item.images
             else:

examples/offline_inference/hunyuan_image3/image_to_text.py

@@ -9,7 +9,7 @@ from PIL import Image
 from vllm_omni.entrypoints.omni import Omni
 
 """
-The tencent/HunyuanImage-3.0-Instruct base model is built on the Hunyuan v1 architecture, specifically the tencent/Hunyuan-A13B-Instruct model. It utilizes two tokenizer delimiter templates:
+The tencent/HunyuanImage-3.0-Instruct base model uses the tencent/Hunyuan-A13B-Instruct backbone. It utilizes two tokenizer delimiter templates:
 
 1) Pretrained template (default for gen_text mode), which concatenates system, image
    tokens, and user question WITHOUT role delimiters:
@@ -70,8 +70,8 @@ def main(args: argparse.Namespace) -> None:
     prompts = [prompt_dict]
     omni_outputs = omni.generate(prompts=prompts)
 
-    prompt_text = omni_outputs[0].request_output[0].prompt
-    generated_text = omni_outputs[0].request_output[0].outputs[0].text
+    prompt_text = omni_outputs[0].request_output.prompt
+    generated_text = omni_outputs[0].request_output.outputs[0].text
     print(f"Prompt: {prompt_text}")
     print(f"Text: {generated_text}")
 

examples/offline_inference/image_to_image/image_edit.py

@@ -463,12 +463,12 @@ def main():
         raise ValueError("No output generated from omni.generate()")
 
     # Extract images from OmniRequestOutput
-    # omni.generate() returns list[OmniRequestOutput], extract images from request_output[0].images
+    # omni.generate() returns list[OmniRequestOutput], extract images from request_output.images
     first_output = outputs[0]
     if not hasattr(first_output, "request_output") or not first_output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")
 
-    req_out = first_output.request_output[0]
+    req_out = first_output.request_output
     if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
         raise ValueError("Invalid request_output structure or missing 'images' key")
 

examples/offline_inference/image_to_video/image_to_video.py

@@ -338,8 +338,6 @@ def main():
             audio = frames.multimodal_output["audio"]
         if frames.is_pipeline_output and frames.request_output is not None:
             inner_output = frames.request_output
-            if isinstance(inner_output, list):
-                inner_output = inner_output[0] if inner_output else None
             if isinstance(inner_output, OmniRequestOutput):
                 if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
                     audio = inner_output.multimodal_output["audio"]

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -99,14 +99,12 @@ def main() -> None:
 
     lines: list[str] = []
     for stage_outputs in outputs:
-        req_outputs = getattr(stage_outputs, "request_output", stage_outputs)
-        req_outputs = req_outputs if isinstance(req_outputs, list) else [req_outputs]
-        for ro in req_outputs:
-            text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
-            lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
-            lines.append("answer:\n")
-            lines.append(text.strip() + "\n")
-            lines.append("\n")
+        ro = getattr(stage_outputs, "request_output", stage_outputs)
+        text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
+        lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
+        lines.append("answer:\n")
+        lines.append(text.strip() + "\n")
+        lines.append("\n")
 
     print("\n".join(lines))
 

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py

@@ -148,19 +148,16 @@ def _collect_images(outputs: list) -> list[torch.Tensor]:
     """Extract all image tensors produced by the final (DiT) stage."""
     images: list[torch.Tensor] = []
     for out in outputs:
-        ro_list = getattr(out, "request_output", out)
-        if not isinstance(ro_list, list):
-            ro_list = [ro_list]
-        for ro_item in ro_list:
-            for completion in getattr(ro_item, "outputs", None) or []:
-                mm = getattr(completion, "multimodal_output", None)
-                if not isinstance(mm, dict) or "image" not in mm:
-                    raise RuntimeError(f"Missing image in multimodal output: {mm}")
-                payload = mm["image"]
-                for tensor in payload if isinstance(payload, list) else [payload]:
-                    if not isinstance(tensor, torch.Tensor):
-                        raise TypeError(f"Expected image tensor, got {type(tensor)}")
-                    images.append(tensor)
+        ro_item = getattr(out, "request_output", out)
+        for completion in getattr(ro_item, "outputs", None) or []:
+            mm = getattr(completion, "multimodal_output", None)
+            if not isinstance(mm, dict) or "image" not in mm:
+                raise RuntimeError(f"Missing image in multimodal output: {mm}")
+            payload = mm["image"]
+            for tensor in payload if isinstance(payload, list) else [payload]:
+                if not isinstance(tensor, torch.Tensor):
+                    raise TypeError(f"Expected image tensor, got {type(tensor)}")
+                images.append(tensor)
     return images
 
 

examples/offline_inference/mimo_audio/end2end.py

@@ -180,11 +180,11 @@ def main(args):
     # Get the query function and call it with appropriate parameters
     query_func = query_map[args.query_type]
 
-    omni_llm = Omni(
+    omni = Omni(
         model=model_name,
         stage_configs_path=args.stage_configs_path,
         log_stats=args.enable_stats,
-        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
+        log_file=("omni_pipeline.log" if args.enable_stats else None),
         init_sleep_seconds=args.init_sleep_seconds,
         batch_timeout=args.batch_timeout,
         init_timeout=args.init_timeout,
@@ -290,7 +290,7 @@ def main(args):
     prompts = [copy.deepcopy(query_result) for _ in range(args.num_prompts)]
 
     print("prompts", prompts)
-    omni_outputs = omni_llm.generate(prompts, sampling_params_list)
+    omni_outputs = omni.generate(prompts, sampling_params_list)
 
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
     if args.query_type is not None:
@@ -298,45 +298,44 @@ def main(args):
     os.makedirs(output_dir, exist_ok=True)
 
     for stage_outputs in omni_outputs:
+        output = stage_outputs.request_output
         if stage_outputs.final_output_type == "text":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                text_output = output.outputs[0].text
-                # Save aligned text file per request
-                prompt_text = output.prompt
-                out_txt = os.path.join(output_dir, f"{request_id}.txt")
-                lines = []
-                lines.append("Prompt:\n")
-                lines.append(str(prompt_text) + "\n")
-                lines.append("vllm_text_output:\n")
-                lines.append(str(text_output).strip() + "\n")
-                try:
-                    with open(out_txt, "w", encoding="utf-8") as f:
-                        print("lines", lines)
-                        f.writelines(lines)
-                except Exception as e:
-                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
-                print(f"Request ID: {request_id}, Text saved to {out_txt}\n")
+            request_id = output.request_id
+            text_output = output.outputs[0].text
+            # Save aligned text file per request
+            prompt_text = output.prompt
+            out_txt = os.path.join(output_dir, f"{request_id}.txt")
+            lines = []
+            lines.append("Prompt:\n")
+            lines.append(str(prompt_text) + "\n")
+            lines.append("vllm_text_output:\n")
+            lines.append(str(text_output).strip() + "\n")
+            try:
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    print("lines", lines)
+                    f.writelines(lines)
+            except Exception as e:
+                print(f"[Warn] Failed writing text file {out_txt}: {e}")
+            print(f"Request ID: {request_id}, Text saved to {out_txt}\n")
         elif stage_outputs.final_output_type == "audio":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                audio_tensor = output.outputs[0].multimodal_output.get("audio")
+            request_id = output.request_id
+            audio_tensor = output.outputs[0].multimodal_output.get("audio")
 
-                if audio_tensor is None:
-                    continue
+            if audio_tensor is None:
+                continue
 
-                output_wav = os.path.join(output_dir, f"{request_id}.wav")
+            output_wav = os.path.join(output_dir, f"{request_id}.wav")
 
-                # Convert to numpy array and ensure correct format
-                audio_numpy = audio_tensor.float().detach().cpu().numpy()
+            # Convert to numpy array and ensure correct format
+            audio_numpy = audio_tensor.float().detach().cpu().numpy()
 
-                # Ensure audio is 1D (flatten if needed)
-                if audio_numpy.ndim > 1:
-                    audio_numpy = audio_numpy.flatten()
+            # Ensure audio is 1D (flatten if needed)
+            if audio_numpy.ndim > 1:
+                audio_numpy = audio_numpy.flatten()
 
-                # Save audio file with explicit WAV format
-                sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
-                print(f"Request ID: {request_id}, Audio saved to {output_wav}")
+            # Save audio file with explicit WAV format
+            sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
+            print(f"Request ID: {request_id}, Audio saved to {output_wav}")
 
 
 def parse_args():

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -320,7 +320,7 @@ def main(args):
         query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate)
     else:
         query_result = query_func()
-    omni_llm = Omni(
+    omni = Omni(
         model=model_name,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
@@ -378,9 +378,11 @@ def main(args):
             prompt["modalities"] = output_modalities
 
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
-    if profiler_enabled:
-        omni_llm.start_profile(stages=[0])
-    omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+    if profiler_enabled and hasattr(omni, "start_profile"):
+        omni.start_profile(stages=[0])
+    elif profiler_enabled:
+        print("[Warn] VLLM_TORCH_PROFILER_DIR is set, but current engine does not support profiler controls.")
+    omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
 
     # Determine output directory: prefer --output-dir; fallback to --output-wav
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
@@ -389,43 +391,42 @@ def main(args):
     total_requests = len(prompts)
     processed_count = 0
     for stage_outputs in omni_generator:
+        output = stage_outputs.request_output
         if stage_outputs.final_output_type == "text":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                text_output = output.outputs[0].text
-                # Save aligned text file per request
-                prompt_text = output.prompt
-                out_txt = os.path.join(output_dir, f"{request_id}.txt")
-                lines = []
-                lines.append("Prompt:\n")
-                lines.append(str(prompt_text) + "\n")
-                lines.append("vllm_text_output:\n")
-                lines.append(str(text_output).strip() + "\n")
-                try:
-                    with open(out_txt, "w", encoding="utf-8") as f:
-                        f.writelines(lines)
-                except Exception as e:
-                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
-                print(f"Request ID: {request_id}, Text saved to {out_txt}")
+            request_id = output.request_id
+            text_output = output.outputs[0].text
+            # Save aligned text file per request
+            prompt_text = output.prompt
+            out_txt = os.path.join(output_dir, f"{request_id}.txt")
+            lines = []
+            lines.append("Prompt:\n")
+            lines.append(str(prompt_text) + "\n")
+            lines.append("vllm_text_output:\n")
+            lines.append(str(text_output).strip() + "\n")
+            try:
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    f.writelines(lines)
+            except Exception as e:
+                print(f"[Warn] Failed writing text file {out_txt}: {e}")
+            print(f"Request ID: {request_id}, Text saved to {out_txt}")
         elif stage_outputs.final_output_type == "audio":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                audio_tensor = output.outputs[0].multimodal_output["audio"]
-                output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
-                sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
-                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+            request_id = output.request_id
+            audio_tensor = output.outputs[0].multimodal_output["audio"]
+            output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
+            sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
+            print(f"Request ID: {request_id}, Saved audio to {output_wav}")
 
-        processed_count += len(stage_outputs.request_output)
-        if profiler_enabled and processed_count >= total_requests:
+        processed_count += 1
+        if profiler_enabled and hasattr(omni, "stop_profile") and processed_count >= total_requests:
             print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
             # Stop the profiler while workers are still alive
-            omni_llm.stop_profile()
+            omni.stop_profile()
 
             print("[Info] Waiting 30s for workers to write massive trace files to disk...")
             time.sleep(30)
             print("[Info] Trace export wait finished.")
 
-    omni_llm.close()
+    omni.close()
 
 
 def parse_args():

examples/offline_inference/qwen3_omni/end2end.py

@@ -294,7 +294,7 @@ def main(args):
     else:
         query_result = query_func()
 
-    omni_llm = Omni(
+    omni = Omni(
         model=model_name,
         stage_configs_path=args.stage_configs_path,
         log_stats=args.log_stats,
@@ -354,8 +354,8 @@ def main(args):
 
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
     if profiler_enabled:
-        omni_llm.start_profile(stages=[0])
-    omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+        omni.start_profile(stages=[0])
+    omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
     # Determine output directory: prefer --output-dir; fallback to --output-wav
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
     os.makedirs(output_dir, exist_ok=True)
@@ -366,51 +366,50 @@ def main(args):
     print(f"query type: {args.query_type}")
 
     for stage_outputs in omni_generator:
+        output = stage_outputs.request_output
         if stage_outputs.final_output_type == "text":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                text_output = output.outputs[0].text
-                # Save aligned text file per request
-                prompt_text = output.prompt
-                out_txt = os.path.join(output_dir, f"{request_id}.txt")
-                lines = []
-                lines.append("Prompt:\n")
-                lines.append(str(prompt_text) + "\n")
-                lines.append("vllm_text_output:\n")
-                lines.append(str(text_output).strip() + "\n")
-                try:
-                    with open(out_txt, "w", encoding="utf-8") as f:
-                        f.writelines(lines)
-                except Exception as e:
-                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
-                print(f"Request ID: {request_id}, Text saved to {out_txt}")
+            request_id = output.request_id
+            text_output = output.outputs[0].text
+            # Save aligned text file per request
+            prompt_text = output.prompt
+            out_txt = os.path.join(output_dir, f"{request_id}.txt")
+            lines = []
+            lines.append("Prompt:\n")
+            lines.append(str(prompt_text) + "\n")
+            lines.append("vllm_text_output:\n")
+            lines.append(str(text_output).strip() + "\n")
+            try:
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    f.writelines(lines)
+            except Exception as e:
+                print(f"[Warn] Failed writing text file {out_txt}: {e}")
+            print(f"Request ID: {request_id}, Text saved to {out_txt}")
         elif stage_outputs.final_output_type == "audio":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                audio_tensor = output.outputs[0].multimodal_output["audio"]
-                output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
+            request_id = output.request_id
+            audio_tensor = output.outputs[0].multimodal_output["audio"]
+            output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
 
-                # Convert to numpy array and ensure correct format
-                audio_numpy = audio_tensor.float().detach().cpu().numpy()
+            # Convert to numpy array and ensure correct format
+            audio_numpy = audio_tensor.float().detach().cpu().numpy()
 
-                # Ensure audio is 1D (flatten if needed)
-                if audio_numpy.ndim > 1:
-                    audio_numpy = audio_numpy.flatten()
+            # Ensure audio is 1D (flatten if needed)
+            if audio_numpy.ndim > 1:
+                audio_numpy = audio_numpy.flatten()
 
-                # Save audio file with explicit WAV format
-                sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
-                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+            # Save audio file with explicit WAV format
+            sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
+            print(f"Request ID: {request_id}, Saved audio to {output_wav}")
 
-        processed_count += len(stage_outputs.request_output)
+        processed_count += 1
         if profiler_enabled and processed_count >= total_requests:
             print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
             # Stop the profiler while workers are still alive
-            omni_llm.stop_profile()
+            omni.stop_profile()
 
             print("[Info] Waiting 30s for workers to write trace files to disk...")
             time.sleep(30)
             print("[Info] Trace export wait time finished.")
-    omni_llm.close()
+    omni.close()
 
 
 def parse_args():

examples/offline_inference/qwen3_omni/end2end_async_chunk.py

@@ -231,52 +231,47 @@ async def run_single_request(
             sampling_params_list=sampling_params_list,
             output_modalities=output_modalities,
         ):
-            if not isinstance(omni_output.request_output, list):
-                outputs_list = [omni_output.request_output]
-            else:
-                outputs_list = omni_output.request_output
+            output = omni_output.request_output
+            if omni_output.final_output_type == "text":
+                if stage_0_first_output_ts is None:
+                    stage_0_first_output_ts = time.perf_counter()
+                text_output = output.outputs[0].text
+                if output.finished:
+                    text_parts.append(text_output)
+            elif omni_output.final_output_type == "audio":
+                mm_out = output.outputs[0].multimodal_output
+                if mm_out and "audio" in mm_out:
+                    if first_audio_ts is None:
+                        first_audio_ts = time.perf_counter()
+                    if audio_sr is None and "sr" in mm_out:
+                        sr_val = mm_out["sr"]
+                        audio_sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
+                        samplerate = audio_sr
+                    audio_data = mm_out["audio"]
+                    if isinstance(audio_data, list):
+                        new_chunks = audio_data[audio_list_consumed:]
+                        audio_list_consumed = len(audio_data)
+                    elif isinstance(audio_data, torch.Tensor):
+                        new_chunks = [audio_data]
+                        audio_last_tensor = audio_data
+                    else:
+                        new_chunks = []
 
-            for output in outputs_list:
-                if omni_output.final_output_type == "text":
-                    if stage_0_first_output_ts is None:
-                        stage_0_first_output_ts = time.perf_counter()
-                    text_output = output.outputs[0].text
-                    if output.finished:
-                        text_parts.append(text_output)
-                elif omni_output.final_output_type == "audio":
-                    mm_out = output.outputs[0].multimodal_output
-                    if mm_out and "audio" in mm_out:
-                        if first_audio_ts is None:
-                            first_audio_ts = time.perf_counter()
-                        if audio_sr is None and "sr" in mm_out:
-                            sr_val = mm_out["sr"]
-                            audio_sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
-                            samplerate = audio_sr
-                        audio_data = mm_out["audio"]
-                        if isinstance(audio_data, list):
-                            new_chunks = audio_data[audio_list_consumed:]
-                            audio_list_consumed = len(audio_data)
-                        elif isinstance(audio_data, torch.Tensor):
-                            new_chunks = [audio_data]
-                            audio_last_tensor = audio_data
-                        else:
-                            new_chunks = []
-
-                        if stream_audio_to_disk and new_chunks:
-                            if sf_writer is None:
-                                sf_writer = sf.SoundFile(
-                                    wav_file,
-                                    mode="w",
-                                    samplerate=samplerate,
-                                    channels=1,
-                                    subtype="FLOAT",
-                                )
-                            for chunk in new_chunks:
-                                chunk_np = chunk.float().detach().cpu().numpy().flatten()
-                                sf_writer.write(chunk_np)
-                                audio_samples_written += len(chunk_np)
-                        else:
-                            audio_chunks.extend(new_chunks)
+                    if stream_audio_to_disk and new_chunks:
+                        if sf_writer is None:
+                            sf_writer = sf.SoundFile(
+                                wav_file,
+                                mode="w",
+                                samplerate=samplerate,
+                                channels=1,
+                                subtype="FLOAT",
+                            )
+                        for chunk in new_chunks:
+                            chunk_np = chunk.float().detach().cpu().numpy().flatten()
+                            sf_writer.write(chunk_np)
+                            audio_samples_written += len(chunk_np)
+                    else:
+                        audio_chunks.extend(new_chunks)
     finally:
         if sf_writer is not None:
             sf_writer.close()

examples/offline_inference/qwen3_tts/end2end.py

@@ -51,16 +51,68 @@ def _estimate_prompt_len(
 
             tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left")
             cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True)
-            _cache[model_name] = (tok, getattr(cfg, "talker_config", None))
 
-        tok, tcfg = _cache[model_name]
+            # Load speech tokenizer (codec encoder) for exact ref_code_len.
+            speech_tok = None
+            try:
+                import os
+
+                from transformers.utils import cached_file
+
+                from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer
+
+                st_cfg_path = cached_file(model_name, "speech_tokenizer/config.json")
+                if st_cfg_path:
+                    speech_tok = Qwen3TTSTokenizer.from_pretrained(
+                        os.path.dirname(st_cfg_path), torch_dtype=torch.bfloat16
+                    )
+                    logger.info("Loaded speech tokenizer for exact ref_code_len estimation")
+            except Exception as e:
+                logger.debug("Could not load speech tokenizer: %s", e)
+
+            _cache[model_name] = (tok, getattr(cfg, "talker_config", None), speech_tok)
+
+        tok, tcfg, speech_tok = _cache[model_name]
         task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]
+
+        def _estimate_ref_code_len(ref_audio: object) -> int | None:
+            """Encode ref_audio with the actual codec to get exact frame count."""
+            if not isinstance(ref_audio, (str, list)):
+                return None
+            audio_path = ref_audio[0] if isinstance(ref_audio, list) else ref_audio
+            if not isinstance(audio_path, str) or not audio_path.strip():
+                return None
+            try:
+                from vllm.multimodal.media import MediaConnector
+
+                connector = MediaConnector(allowed_local_media_path="/")
+                audio, sr = connector.fetch_audio(audio_path)
+                import numpy as np
+
+                wav_np = np.asarray(audio, dtype=np.float32)
+
+                if speech_tok is not None:
+                    enc = speech_tok.encode(wav_np, sr=int(sr), return_dict=True)
+                    ref_code = getattr(enc, "audio_codes", None)
+                    if isinstance(ref_code, list):
+                        ref_code = ref_code[0] if ref_code else None
+                    if ref_code is not None and hasattr(ref_code, "shape"):
+                        shape = ref_code.shape
+                        return int(shape[0]) if len(shape) == 2 else int(shape[1]) if len(shape) == 3 else None
+
+                # Fallback: estimate from duration
+                codec_hz = getattr(tcfg, "codec_frame_rate", None) or 12
+                return int(len(audio) / sr * codec_hz)
+            except Exception:
+                return None
+
         return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
             additional_information=additional_information,
             task_type=task_type,
             tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
             codec_language_id=getattr(tcfg, "codec_language_id", None),
             spk_is_dialect=getattr(tcfg, "spk_is_dialect", None),
+            estimate_ref_code_len=_estimate_ref_code_len,
         )
     except Exception as exc:
         logger.warning("Failed to estimate prompt length, using fallback 2048: %s", exc)
@@ -325,8 +377,8 @@ def main(args):
     for batch_start in range(0, len(inputs), batch_size):
         batch = inputs[batch_start : batch_start + batch_size]
         for stage_outputs in omni.generate(batch):
-            for output in stage_outputs.request_output:
-                _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
+            output = stage_outputs.request_output
+            _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
 
 
 async def main_streaming(args):

examples/offline_inference/text_to_audio/text_to_audio.py

@@ -177,7 +177,7 @@ def main():
     output = outputs[0]
     if not hasattr(output, "request_output") or not output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")
-    request_output = output.request_output[0]
+    request_output = output.request_output
     if not hasattr(request_output, "multimodal_output"):
         raise ValueError("No multimodal_output found in request_output")
 

examples/offline_inference/text_to_image/README.md

@@ -16,7 +16,7 @@ if __name__ == "__main__":
     omni = Omni(model="Qwen/Qwen-Image")
     prompt = "a cup of coffee on the table"
     outputs = omni.generate(prompt)
-    images = outputs[0].request_output[0].images
+    images = outputs[0].request_output.images
     images[0].save("coffee.png")
 ```
 
@@ -34,7 +34,7 @@ if __name__ == "__main__":
     ]
     outputs = omni.generate(prompts)
     for i, output in enumerate(outputs):
-      image = output.request_output[0].images[0].save(f"{i}.jpg")
+      image = output.request_output.images[0].save(f"{i}.jpg")
 ```
 
 !!! info
@@ -69,7 +69,7 @@ if __name__ == "__main__":
       }
     ])
     for i, output in enumerate(outputs):
-      image = output.request_output[0].images[0].save(f"{i}.jpg")
+      image = output.request_output.images[0].save(f"{i}.jpg")
 ```
 
 ## Local CLI Usage

examples/offline_inference/text_to_image/gradio_demo.py

@@ -112,7 +112,7 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
         )
         images_outputs = []
         for output in outputs:
-            req_out = output.request_output[0]
+            req_out = output.request_output
             if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
                 raise ValueError("Invalid request_output structure or missing 'images' key")
             images = req_out.images

examples/offline_inference/text_to_image/text_to_image.py

@@ -14,7 +14,6 @@ from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
-from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
 
 
@@ -400,20 +399,18 @@ def main():
         else:
             print("[Profiler] No valid profiling data returned.")
 
-    # Extract images from OmniRequestOutput
-    # omni.generate() returns list[OmniRequestOutput], extract images from the first output
+    # omni.generate() returns list[OmniRequestOutput]
     if not outputs or len(outputs) == 0:
         raise ValueError("No output generated from omni.generate()")
     logger.info(f"Outputs: {outputs}")
 
-    # Extract images from request_output[0]['images']
     first_output = outputs[0]
     if not hasattr(first_output, "request_output") or not first_output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")
 
-    req_out = first_output.request_output[0]
-    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
-        raise ValueError("Invalid request_output structure or missing 'images' key")
+    req_out = first_output.request_output
+    if not hasattr(req_out, "images"):
+        raise ValueError("Invalid request_output structure or missing 'images'.")
 
     images = req_out.images
     if not images:

examples/offline_inference/text_to_video/text_to_video.py

@@ -240,8 +240,6 @@ def main():
             audio = frames.multimodal_output["audio"]
         if frames.is_pipeline_output and frames.request_output is not None:
             inner_output = frames.request_output
-            if isinstance(inner_output, list):
-                inner_output = inner_output[0] if inner_output else None
             if isinstance(inner_output, OmniRequestOutput):
                 if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
                     audio = inner_output.multimodal_output["audio"]

requirements/common.txt

@@ -17,3 +17,4 @@ einops>=0.8.1
 prettytable>=3.8.0
 aenum==3.1.16
 pyzmq>=25.0.0
+janus>=1.0.0

tests/comfyui/test_comfyui_integration.py

@@ -11,6 +11,7 @@ import time
 import traceback
 from collections.abc import Iterable, Sequence
 from enum import StrEnum, auto
+from types import SimpleNamespace
 from typing import Any, NamedTuple
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -31,6 +32,7 @@ from vllm import SamplingParams
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
+from vllm_omni.entrypoints.async_omni import AsyncOmni as RealAsyncOmni
 from vllm_omni.entrypoints.cli.serve import OmniServeCommand
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -43,7 +45,7 @@ class ServerCase(NamedTuple):
 
     served_model: str
     stage_list: list
-    stage_configs: list[dict]
+    stage_configs: list[Any]
     outputs: list[OmniRequestOutput]
 
 
@@ -262,6 +264,40 @@ def _assert_model_param_values(received: OmniSamplingParams, expected: dict):
         )
 
 
+def _make_stage_config(
+    stage_type: str,
+    *,
+    is_comprehension: bool = False,
+    model_stage: str | None = None,
+):
+    engine_args = SimpleNamespace()
+    if model_stage is not None:
+        engine_args.model_stage = model_stage
+    return SimpleNamespace(
+        stage_type=stage_type,
+        is_comprehension=is_comprehension,
+        engine_args=engine_args,
+    )
+
+
+def _stage_type(stage: Any) -> str | None:
+    return getattr(stage, "stage_type", None)
+
+
+def _build_output_modalities(stage_configs: list[Any]) -> list[str]:
+    modalities: list[str] = []
+    for stage in stage_configs:
+        final_output = getattr(stage, "final_output", False)
+        final_output_type = getattr(stage, "final_output_type", None)
+        if final_output and isinstance(final_output_type, str):
+            modalities.append(final_output_type)
+    if modalities:
+        return modalities
+    if any(_stage_type(stage) == "diffusion" for stage in stage_configs):
+        return ["image"]
+    return ["text", "audio"]
+
+
 def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: SamplingCase, server_case: ServerCase):
     async def _mock_generate(*args, **kwargs):
         received_sampling_params_list: Sequence[OmniSamplingParams] | None = (
@@ -367,18 +403,28 @@ def mock_async_omni(server_case: ServerCase, sampling_case: SamplingCase):
             new=_mock_preprocess_chat,
         ),
     ):
-        mock_instance = AsyncMock()
+        mock_instance = AsyncMock(spec=RealAsyncOmni)
         mock_instance.generate = _build_mock_outputs(server_case.outputs, sampling_case, server_case)
 
         mock_instance.stage_list = server_case.stage_list
         mock_instance.stage_configs = server_case.stage_configs
+        mock_instance.output_modalities = _build_output_modalities(server_case.stage_configs)
         mock_instance.default_sampling_params_list = [
-            SamplingParams() if stage.get("stage_type") != "diffusion" else MagicMock()
+            SamplingParams() if _stage_type(stage) != "diffusion" else MagicMock()
             for stage in server_case.stage_configs
         ]
         mock_instance.errored = False
         mock_instance.dead_error = RuntimeError("Mock engine error")
-        mock_instance.model_config = MagicMock(max_model_len=4096, io_processor_plugin=None)
+        mock_instance.model_config = MagicMock(
+            max_model_len=4096,
+            io_processor_plugin=None,
+            allowed_local_media_path=None,
+            allowed_media_domains=None,
+        )
+        # Mimic Qwen3-TTS talker speaker config so CustomVoice validation passes.
+        mock_instance.model_config.hf_config = MagicMock()
+        mock_instance.model_config.hf_config.talker_config = MagicMock()
+        mock_instance.model_config.hf_config.talker_config.speaker_id = {"Vivian": 0}
         mock_instance.io_processor = MagicMock()
         mock_instance.input_processor = MagicMock()
         mock_instance.shutdown = MagicMock()
@@ -444,7 +490,7 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
             ServerCase(
                 served_model="Tongyi-MAI/Z-Image-Turbo",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion"}],
+                stage_configs=[_make_stage_config("diffusion")],
                 outputs=[_build_diffusion_image_output_for_images_endpoint()],
             ),
             "Tongyi-MAI/Z-Image-Turbo",
@@ -455,7 +501,7 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
             ServerCase(
                 served_model="ByteDance-Seed/BAGEL-7B-MoT",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion"}],
+                stage_configs=[_make_stage_config("diffusion")],
                 outputs=[_build_diffusion_image_output_for_chat_endpoint()],
             ),
             "ByteDance-Seed/BAGEL-7B-MoT",
@@ -466,7 +512,7 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
             ServerCase(
                 served_model="Qwen/Qwen-Image-Edit",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion"}],
+                stage_configs=[_make_stage_config("diffusion")],
                 outputs=[_build_diffusion_image_output_for_images_endpoint()],
             ),
             "Qwen/Qwen-Image-Edit",
@@ -477,7 +523,7 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
             ServerCase(
                 served_model="ByteDance-Seed/BAGEL-7B-MoT",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion"}],
+                stage_configs=[_make_stage_config("diffusion")],
                 outputs=[_build_diffusion_image_output_for_chat_endpoint()],
             ),
             "ByteDance-Seed/BAGEL-7B-MoT",
@@ -542,9 +588,9 @@ async def test_image_generation_node(api_server: str, model: str, image_input: b
                     MagicMock(is_comprehension=False, model_stage="llm"),
                 ],
                 stage_configs=[
-                    {"stage_type": "llm"},
-                    {"stage_type": "llm"},
-                    {"stage_type": "llm"},
+                    _make_stage_config("llm", is_comprehension=True, model_stage="thinker"),
+                    _make_stage_config("llm", is_comprehension=False, model_stage="talker"),
+                    _make_stage_config("llm", is_comprehension=False, model_stage="code2wav"),
                 ],
                 outputs=[_build_audio_chat_output(), _build_text_output("Understanding response")],
             ),
@@ -599,7 +645,7 @@ async def test_understanding_node(api_server: str, sampling_case: SamplingCase):
             ServerCase(
                 served_model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
                 stage_list=["llm"],
-                stage_configs=[{"stage_type": "llm"}],
+                stage_configs=[_make_stage_config("llm", model_stage="qwen3_tts")],
                 outputs=[_build_audio_speech_output()],
             ),
             VLLMOmniTTS,
@@ -617,7 +663,7 @@ async def test_understanding_node(api_server: str, sampling_case: SamplingCase):
             ServerCase(
                 served_model="Qwen/Qwen3-TTS-12Hz-1.7B-Base",
                 stage_list=["llm"],
-                stage_configs=[{"stage_type": "llm"}],
+                stage_configs=[_make_stage_config("llm", model_stage="qwen3_tts")],
                 outputs=[_build_audio_speech_output()],
             ),
             VLLMOmniVoiceClone,

tests/conftest.py

@@ -1733,7 +1733,9 @@ class OmniRunner:
         Returns:
             List of SamplingParams with default decoding for each stage
         """
-        return [st.default_sampling_params for st in self.omni.stage_list]
+        if not hasattr(self.omni, "default_sampling_params_list"):
+            raise AttributeError("Omni.default_sampling_params_list is not available")
+        return list(self.omni.default_sampling_params_list)
 
     def get_omni_inputs(
         self,
@@ -1974,9 +1976,9 @@ class OmniRunnerHandler:
             audio_content = None
             for stage_output in outputs:
                 if getattr(stage_output, "final_output_type", None) == "text":
-                    text_content = stage_output.request_output[0].outputs[0].text
+                    text_content = stage_output.request_output.outputs[0].text
                 if getattr(stage_output, "final_output_type", None) == "audio":
-                    audio_content = stage_output.request_output[0].outputs[0].multimodal_output["audio"]
+                    audio_content = stage_output.request_output.outputs[0].multimodal_output["audio"]
 
             result.audio_content = audio_content
             result.text_content = text_content

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -70,6 +70,7 @@ def mock_dist():
         patch.object(dist, "get_world_size", return_value=2),
         patch.object(dist, "get_rank", return_value=0),
         patch.object(dist, "is_initialized", return_value=True),
+        patch.object(dist, "all_reduce", return_value=None),
         patch.object(dist, "gather", return_value=None),
         patch.object(dist, "broadcast", return_value=None),
     ):

tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml

@@ -1,4 +1,4 @@
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni for multi-stage omni runtime.
 
 # This config is optimized for CI e2e tests.
 stage_args:

tests/e2e/offline_inference/test_bagel_img2img.py

@@ -98,9 +98,9 @@ def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
         if images := getattr(req_output, "images", None):
             return images[0]
         if hasattr(req_output, "request_output") and req_output.request_output:
-            for stage_out in req_output.request_output:
-                if hasattr(stage_out, "images") and stage_out.images:
-                    return stage_out.images[0]
+            stage_out = req_output.request_output
+            if hasattr(stage_out, "images") and stage_out.images:
+                return stage_out.images[0]
     return None
 
 

tests/e2e/offline_inference/test_bagel_text2img.py

@@ -14,6 +14,9 @@ Equivalent to running:
 """
 
 import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 import signal
 import socket
 import subprocess
@@ -96,9 +99,9 @@ def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
         if images := getattr(req_output, "images", None):
             return images[0]
         if hasattr(req_output, "request_output") and req_output.request_output:
-            for stage_out in req_output.request_output:
-                if hasattr(stage_out, "images") and stage_out.images:
-                    return stage_out.images[0]
+            stage_out = req_output.request_output
+            if hasattr(stage_out, "images") and stage_out.images:
+                return stage_out.images[0]
     return None
 
 

tests/e2e/offline_inference/test_cache_dit.py

@@ -72,13 +72,13 @@ def test_cache_dit(model_name: str):
                 num_outputs_per_prompt=1,  # Single output for speed
             ),
         )
-        # Extract images from request_output[0]['images']
+        # Extract images from request_output['images']
         first_output = outputs[0]
         assert first_output.final_output_type == "image"
         if not hasattr(first_output, "request_output") or not first_output.request_output:
             raise ValueError("No request_output found in OmniRequestOutput")
 
-        req_out = first_output.request_output[0]
+        req_out = first_output.request_output
         if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
             raise ValueError("Invalid request_output structure or missing 'images' key")
 

tests/e2e/offline_inference/test_diffusion_lora.py

@@ -36,7 +36,7 @@ def test_diffusion_model(model_name: str, tmp_path: Path):
         if not hasattr(first_output, "request_output") or not first_output.request_output:
             raise ValueError("No request_output found in OmniRequestOutput")
 
-        req_out = first_output.request_output[0]
+        req_out = first_output.request_output
         if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
             raise ValueError("Invalid request_output structure or missing 'images' key")
         return req_out.images

tests/e2e/offline_inference/test_expert_parallel.py

@@ -129,7 +129,7 @@ def _run_inference(
         elapsed_ms = (time.time() - start) * 1000
 
         return InferenceResult(
-            images=outputs[0].request_output[0].images,
+            images=outputs[0].images,
             elapsed_ms=elapsed_ms,
         )
     finally:

tests/e2e/offline_inference/test_sequence_parallel.py

@@ -130,7 +130,7 @@ def _run_inference(
         elapsed_ms = (time.time() - start) * 1000
 
         return InferenceResult(
-            images=outputs[0].request_output[0].images,
+            images=outputs[0].request_output.images,
             elapsed_ms=elapsed_ms,
         )
     finally:

tests/e2e/offline_inference/test_stable_audio_model.py

@@ -57,7 +57,7 @@ def test_stable_audio_model(model_name: str):
     assert first_output.final_output_type == "image"
     assert hasattr(first_output, "request_output") and first_output.request_output
 
-    req_out = first_output.request_output[0]
+    req_out = first_output.request_output
     assert isinstance(req_out, OmniRequestOutput)
     assert req_out.final_output_type == "audio"
     assert hasattr(req_out, "multimodal_output") and req_out.multimodal_output

tests/e2e/offline_inference/test_t2i_model.py

@@ -28,7 +28,7 @@ if current_omni_platform.is_npu():
     models = ["Tongyi-MAI/Z-Image-Turbo", "Qwen/Qwen-Image"]
 elif current_omni_platform.is_rocm():
     # TODO: When ROCm support is ready, remove this branch.
-    # vLLM V0.11.0 has issues running riverclouds/qwen_image_random
+    # Current upstream vLLM has issues running riverclouds/qwen_image_random
     # on ROCm
     models = ["Tongyi-MAI/Z-Image-Turbo"]
 
@@ -62,13 +62,13 @@ def test_diffusion_model(model_name: str, run_level):
                 num_outputs_per_prompt=2,
             ),
         )
-        # Extract images from request_output[0]['images']
+        # Extract images from request_output['images']
         first_output = outputs[0]
         assert first_output.final_output_type == "image"
         if not hasattr(first_output, "request_output") or not first_output.request_output:
             raise ValueError("No request_output found in OmniRequestOutput")
 
-        req_out = first_output.request_output[0]
+        req_out = first_output.request_output
         if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
             raise ValueError("Invalid request_output structure or missing 'images' key")
 

tests/e2e/offline_inference/test_t2v_model.py

@@ -51,7 +51,7 @@ def test_video_diffusion_model(model_name: str):
     if not hasattr(first_output, "request_output") or not first_output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")
 
-    req_out = first_output.request_output[0]
+    req_out = first_output.request_output
     if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
         raise ValueError("Invalid request_output structure or missing 'images' key")
 

tests/e2e/offline_inference/test_teacache.py

@@ -68,13 +68,13 @@ def test_teacache(model_name: str):
                 num_outputs_per_prompt=1,  # Single output for speed
             ),
         )
-        # Extract images from request_output[0]['images']
+        # Extract images from request_output['images']
         first_output = outputs[0]
         assert first_output.final_output_type == "image"
         if not hasattr(first_output, "request_output") or not first_output.request_output:
             raise ValueError("No request_output found in OmniRequestOutput")
 
-        req_out = first_output.request_output[0]
+        req_out = first_output.request_output
         if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
             raise ValueError("Invalid request_output structure or missing 'images' key")
 

tests/e2e/offline_inference/test_vae_decode_parallelism.py

@@ -72,6 +72,7 @@ def is_nextstep_model(model_name: str) -> bool:
 
 
 def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile, vae_patch_parallel_size=1):
+    m = None
     try:
         parallel_config = DiffusionParallelConfig(
             tensor_parallel_size=tp,
@@ -105,7 +106,7 @@ def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile,
         )
         end = time.perf_counter()
         first_output = outputs[0]
-        req_out = first_output.request_output[0]
+        req_out = first_output.request_output
         frames = req_out.images[0]
         if isinstance(frames, torch.Tensor):
             frames = frames.detach().cpu().numpy()
@@ -115,7 +116,8 @@ def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile,
         cost = (end - start) * 1000
         return frames, cost
     finally:
-        m.close()
+        if m is not None:
+            m.close()
         cleanup_dist_env_and_memory()
 
 

tests/e2e/offline_inference/test_zimage_parallelism.py

@@ -68,7 +68,7 @@ def _extract_single_image(outputs) -> Image.Image:
     if not hasattr(first_output, "request_output") or not first_output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")
 
-    req_out = first_output.request_output[0]
+    req_out = first_output.request_output
     if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
         raise ValueError("Invalid request_output structure or missing 'images' key")
 

tests/e2e/stage_configs/qwen2_5_omni_ci.yaml

@@ -1,4 +1,4 @@
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni for multi-stage omni runtime.
 
 # The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090).
 # This config is optimized for CI e2e tests.

tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml

@@ -1,4 +1,4 @@
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni for multi-stage omni runtime.
 
 # The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090).
 # This config is optimized for CI e2e tests.

tests/e2e/stage_configs/xpu/qwen2_5_omni_ci.yaml

@@ -1,9 +1,9 @@
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni for multi-stage omni runtime.
 
 # The following config is verified with 2 * Intel Arc Pro B60 XPU.
 stage_args:
   - stage_id: 0
-    stage_type: llm # Use llm stage type to launch OmniLLM
+    stage_type: llm # Use llm stage type for AR stages
     runtime:
       process: true # Run this stage in a separate process
       devices: "0" # Visible devices for this stage
@@ -35,7 +35,7 @@ stage_args:
       detokenize: True
       repetition_penalty: 1.1
   - stage_id: 1
-    stage_type: llm # Use llm stage type to launch OmniLLM
+    stage_type: llm # Use llm stage type for AR stages
     runtime:
       process: true
       devices: "1"
@@ -67,7 +67,7 @@ stage_args:
       stop_token_ids: [8294]
 
   - stage_id: 2
-    stage_type: llm # Use llm stage type to launch OmniLLM
+    stage_type: llm # Use llm stage type for AR stages
     runtime:
       process: true
       devices: "2"

tests/e2e/stage_configs/xpu/qwen3_omni_ci.yaml

@@ -6,7 +6,7 @@
 # The following config is verified with 8 * Intel Arc Pro B60 XPU.
 stage_args:
 - stage_id: 0
-  stage_type: llm # Use llm stage type to launch OmniLLM
+  stage_type: llm # Use llm stage type for AR stages
   runtime:
     devices: "0,1,2,3"
     max_batch_size: 1
@@ -41,7 +41,7 @@ stage_args:
     repetition_penalty: 1.05
 
 - stage_id: 1
-  stage_type: llm # Use llm stage type to launch OmniLLM
+  stage_type: llm # Use llm stage type for AR stages
   runtime:
     devices: "4"
     max_batch_size: 1
@@ -75,7 +75,7 @@ stage_args:
     stop_token_ids: [2150]
 
 - stage_id: 2
-  stage_type: llm # Use llm stage type to launch OmniLLM
+  stage_type: llm # Use llm stage type for AR stages
   runtime:
     devices: "5"
     max_batch_size: 1

tests/engine/test_async_omni_engine_input.py

@@ -0,0 +1,63 @@
+from unittest.mock import Mock
+
+import pytest
+from vllm.sampling_params import SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+
+from vllm_omni.engine import OmniEngineCoreRequest
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_engine_core_request() -> EngineCoreRequest:
+    return EngineCoreRequest(
+        request_id="req-1",
+        prompt_token_ids=[1, 1, 1],
+        mm_features=None,
+        sampling_params=SamplingParams(max_tokens=8),
+        pooling_params=None,
+        arrival_time=0.0,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+    )
+
+
+def test_build_add_request_message_preserves_additional_information():
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.default_sampling_params_list = [params]
+    engine.stage_metadata = [{"stage_type": "llm"}]
+    engine.supported_tasks = ("speech",)
+
+    input_processor = Mock()
+    input_processor.process_inputs.return_value = _make_engine_core_request()
+    engine.input_processor = input_processor
+
+    output_processor = Mock()
+    engine.output_processors = [output_processor]
+
+    prompt = {
+        "prompt_token_ids": [1, 1, 1],
+        "additional_information": {
+            "text": ["hello world"],
+            "speaker": ["vivian"],
+        },
+    }
+
+    msg = engine._build_add_request_message(
+        request_id="req-1",
+        prompt=prompt,
+        sampling_params_list=[params],
+        final_stage_id=0,
+        arrival_time=0.0,
+    )
+
+    request = msg["prompt"]
+    assert isinstance(request, OmniEngineCoreRequest)
+    assert request.external_req_id == "req-1"
+    assert request.additional_information is not None
+    assert request.additional_information.entries["text"].list_data == ["hello world"]
+    assert request.additional_information.entries["speaker"].list_data == ["vivian"]
+    output_processor.add_request.assert_called_once()

tests/entrypoints/openai_api/test_image_server.py

@@ -10,6 +10,7 @@ OpenAI-compatible async text-to-image generation API endpoints in api_server.py.
 import base64
 import io
 from argparse import Namespace
+from types import SimpleNamespace
 
 import pytest
 from fastapi.testclient import TestClient
@@ -115,7 +116,10 @@ class FakeAsyncOmni:
     """Fake AsyncOmni that yields a single diffusion output."""
 
     def __init__(self):
-        self.stage_list = ["llm", "diffusion"]
+        self.stage_configs = [
+            SimpleNamespace(stage_type="llm"),
+            SimpleNamespace(stage_type="diffusion"),
+        ]
         self.default_sampling_params_list = [SamplingParams(temperature=0.1), OmniDiffusionSamplingParams()]
         self.captured_sampling_params_list = None
         self.captured_prompt = None
@@ -129,21 +133,25 @@ class FakeAsyncOmni:
 
 @pytest.fixture
 def mock_async_diffusion(mocker: MockerFixture):
-    """Mock AsyncOmniDiffusion instance that returns fake images"""
-    mock = mocker.Mock()
-    mock.is_running = True  # For health endpoint
-    mock.check_health = mocker.AsyncMock()  # For LLM mode health check
+    """Mock diffusion engine that matches the current async-generator API."""
 
-    async def generate(**kwargs):
-        # Return n PIL images wrapped in result object
-        n = kwargs["sampling_params_list"][0].num_outputs_per_prompt
-        mock.captured_sampling_params_list = kwargs["sampling_params_list"]
-        mock.captured_prompt = kwargs["prompt"]
-        images = [Image.new("RGB", (64, 64), color="blue") for _ in range(n)]
-        return MockGenerationResult(images)
+    class MockAsyncDiffusion:
+        def __init__(self) -> None:
+            self.is_running = True
+            self.check_health = mocker.AsyncMock()
+            self.captured_sampling_params_list = None
+            self.captured_prompt = None
+            self.generate_calls = 0
 
-    mock.generate = mocker.AsyncMock(side_effect=generate)
-    return mock
+        async def generate(self, **kwargs):
+            self.generate_calls += 1
+            n = kwargs["sampling_params_list"][0].num_outputs_per_prompt
+            self.captured_sampling_params_list = kwargs["sampling_params_list"]
+            self.captured_prompt = kwargs["prompt"]
+            images = [Image.new("RGB", (64, 64), color="blue") for _ in range(n)]
+            yield MockGenerationResult(images)
+
+    return MockAsyncDiffusion()
 
 
 @pytest.fixture
@@ -159,7 +167,7 @@ def test_client(mock_async_diffusion):
     # Set up app state with diffusion engine
     app.state.engine_client = mock_async_diffusion
     app.state.diffusion_engine = mock_async_diffusion  # Also set for health endpoint
-    app.state.stage_configs = [{"stage_type": "diffusion"}]
+    app.state.stage_configs = [SimpleNamespace(stage_type="diffusion")]
     app.state.diffusion_model_name = "Qwen/Qwen-Image"  # For models endpoint
     app.state.args = Namespace(
         default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5}}',
@@ -180,7 +188,32 @@ def async_omni_test_client():
     app.include_router(router)
 
     app.state.engine_client = FakeAsyncOmni()
-    app.state.stage_configs = [{"stage_type": "llm"}, {"stage_type": "diffusion"}]
+    app.state.stage_configs = [
+        SimpleNamespace(stage_type="llm"),
+        SimpleNamespace(stage_type="diffusion"),
+    ]
+    app.state.args = Namespace(
+        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
+        max_generated_image_size=4096,  # 64*64
+    )
+    return TestClient(app)
+
+
+@pytest.fixture
+def async_omni_stage_configs_only_client():
+    """Create test client with refactored AsyncOmni compatibility surface only."""
+    from fastapi import FastAPI
+
+    from vllm_omni.entrypoints.openai.api_server import router
+
+    app = FastAPI()
+    app.include_router(router)
+
+    engine = FakeAsyncOmni()
+    assert not hasattr(engine, "stage_list")
+    app.state.engine_client = engine
+    # Intentionally do not populate app.state.stage_configs. Refactored
+    # AsyncOmni exposes stage_configs on the engine instance.
     app.state.args = Namespace(
         default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
         max_generated_image_size=4096,  # 64*64
@@ -291,6 +324,45 @@ def test_generate_images_async_omni_sampling_params(async_omni_test_client):
     assert captured[1].seed == 7
 
 
+def test_generate_images_async_omni_stage_configs_only(async_omni_stage_configs_only_client):
+    """Regression: image generation accepts refactored AsyncOmni without stage_list."""
+    response = async_omni_stage_configs_only_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a castle",
+            "n": 1,
+            "size": "256x256",
+            "seed": 11,
+        },
+    )
+    assert response.status_code == 200
+    data = response.json()
+    assert len(data["data"]) == 1
+    engine = async_omni_stage_configs_only_client.app.state.engine_client
+    captured = engine.captured_sampling_params_list
+    assert captured is not None
+    assert len(captured) == 2
+    assert captured[1].seed == 11
+
+
+def test_image_edits_async_omni_stage_configs_only(async_omni_stage_configs_only_client):
+    """Regression: image edits accepts refactored AsyncOmni without stage_list."""
+    img_bytes = make_test_image_bytes((16, 16))
+    response = async_omni_stage_configs_only_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "size": "auto",
+        },
+    )
+    assert response.status_code == 200
+    engine = async_omni_stage_configs_only_client.app.state.engine_client
+    captured = engine.captured_sampling_params_list
+    assert captured is not None
+    assert len(captured) == 2
+
+
 def test_generate_multiple_images(test_client):
     """Test generating multiple images"""
     response = test_client.post(
@@ -538,13 +610,12 @@ def test_parameters_passed_through(test_client, mock_async_diffusion):
     )
     assert response.status_code == 200
 
-    # Ensure generate() was called exactly once
-    mock_async_diffusion.generate.assert_awaited_once()
-    call_kwargs = mock_async_diffusion.generate.call_args[1]["sampling_params_list"][0]
-    assert call_kwargs.num_inference_steps == 100
-    assert call_kwargs.guidance_scale == 7.5
-    assert call_kwargs.true_cfg_scale == 3.0
-    assert call_kwargs.seed == 42
+    assert mock_async_diffusion.generate_calls == 1
+    captured = mock_async_diffusion.captured_sampling_params_list[0]
+    assert captured.num_inference_steps == 100
+    assert captured.guidance_scale == 7.5
+    assert captured.true_cfg_scale == 3.0
+    assert captured.seed == 42
 
 
 def test_model_field_omitted_works(test_client):

tests/entrypoints/openai_api/test_serving_chat_sampling_params.py

@@ -63,9 +63,9 @@ def mock_engine_client(
     default_other_params,
     mocker: MockerFixture,
 ):
-    """Create mock engine client with stage_list and default_sampling_params_list."""
+    """Create mock engine client with stage_configs and default_sampling_params_list."""
     engine_client = mocker.MagicMock()
-    engine_client.stage_list = [mock_comprehension_stage, mock_other_stage]
+    engine_client.stage_configs = [mock_comprehension_stage, mock_other_stage]
     engine_client.default_sampling_params_list = [
         default_comprehension_params,
         default_other_params,
@@ -311,7 +311,7 @@ def test_get_comprehension_stage_index_finds_second_stage(mocker: MockerFixture)
     comprehension.is_comprehension = True
 
     instance.engine_client = mocker.MagicMock()
-    instance.engine_client.stage_list = [other, comprehension]
+    instance.engine_client.stage_configs = [other, comprehension]
 
     assert instance._get_comprehension_stage_index() == 1
 
@@ -328,7 +328,7 @@ def test_get_comprehension_stage_index_raises_when_not_found(mocker: MockerFixtu
     stage2.is_comprehension = False
 
     instance.engine_client = mocker.MagicMock()
-    instance.engine_client.stage_list = [stage1, stage2]
+    instance.engine_client.stage_configs = [stage1, stage2]
 
     with pytest.raises(ValueError, match="No comprehension stage"):
         instance._get_comprehension_stage_index()

tests/entrypoints/openai_api/test_serving_speech.py

@@ -487,7 +487,7 @@ class TestTTSMethods:
     def speech_server(self, mocker: MockerFixture):
         mock_engine_client = mocker.MagicMock()
         mock_engine_client.errored = False
-        mock_engine_client.stage_list = None
+        mock_engine_client.stage_configs = []
         mock_engine_client.tts_max_instructions_length = None
         mock_models = mocker.MagicMock()
         mock_models.is_base_model.return_value = True
@@ -499,7 +499,7 @@ class TestTTSMethods:
 
     def test_is_tts_detection_no_stage(self, speech_server):
         """Test TTS model detection when no TTS stage exists."""
-        # Fixture creates server with stage_list = None -> _is_tts should be False
+        # Fixture creates server with stage_configs = [] -> _is_tts should be False
         assert speech_server._is_tts is False
         assert speech_server._tts_stage is None
 
@@ -511,9 +511,9 @@ class TestTTSMethods:
 
         # Create a TTS stage
         mock_stage = mocker.MagicMock()
-        mock_stage.model_stage = "qwen3_tts"
+        mock_stage.engine_args.model_stage = "qwen3_tts"
         mock_stage.tts_args = {}
-        mock_engine_client.stage_list = [mock_stage]
+        mock_engine_client.stage_configs = [mock_stage]
 
         mock_models = mocker.MagicMock()
         mock_models.is_base_model.return_value = True
@@ -614,7 +614,7 @@ class TestTTSMethods:
         """Test _load_supported_speakers."""
         mock_engine_client = mocker.MagicMock()
         mock_engine_client.errored = False
-        mock_engine_client.stage_list = None
+        mock_engine_client.stage_configs = []
 
         # Mock talker_config with mixed-case speaker names
         mock_talker_config = mocker.MagicMock()
@@ -757,7 +757,7 @@ class TestTTSMethods:
         """Test CLI override (stored in engine_client) takes highest priority."""
         mock_engine_client = mocker.MagicMock()
         mock_engine_client.errored = False
-        mock_engine_client.stage_list = None
+        mock_engine_client.stage_configs = []
         # CLI override is stored in engine_client
         mock_engine_client.tts_max_instructions_length = 1000
         mock_models = mocker.MagicMock()
@@ -781,9 +781,9 @@ class TestTTSMethods:
 
         # Mock stage with tts_args
         mock_stage = mocker.MagicMock()
-        mock_stage.model_stage = "qwen3_tts"
+        mock_stage.engine_args.model_stage = "qwen3_tts"
         mock_stage.tts_args = {"max_instructions_length": 750}
-        mock_engine_client.stage_list = [mock_stage]
+        mock_engine_client.stage_configs = [mock_stage]
 
         server = OmniOpenAIServingSpeech(
             engine_client=mock_engine_client,
@@ -804,9 +804,9 @@ class TestTTSMethods:
 
         # Mock stage with tts_args
         mock_stage = mocker.MagicMock()
-        mock_stage.model_stage = "qwen3_tts"
+        mock_stage.engine_args.model_stage = "qwen3_tts"
         mock_stage.tts_args = {"max_instructions_length": 750}
-        mock_engine_client.stage_list = [mock_stage]
+        mock_engine_client.stage_configs = [mock_stage]
 
         server = OmniOpenAIServingSpeech(
             engine_client=mock_engine_client,
@@ -820,7 +820,7 @@ class TestTTSMethods:
         """Test instructions length validation uses cached _max_instructions_length."""
         mock_engine_client = mocker.MagicMock()
         mock_engine_client.errored = False
-        mock_engine_client.stage_list = None
+        mock_engine_client.stage_configs = []
         # CLI override with max length of 10 characters
         mock_engine_client.tts_max_instructions_length = 10
         mock_models = mocker.MagicMock()
@@ -1034,15 +1034,12 @@ class TestAsyncOmniSupportedTasks:
     @pytest.mark.asyncio
     async def test_tts_only_no_generate_task(self):
         """TTS-only models (audio output, no text) should not include 'generate'."""
-        from unittest.mock import MagicMock
+        from types import SimpleNamespace
 
         from vllm_omni.entrypoints.async_omni import AsyncOmni
 
         omni = AsyncOmni.__new__(AsyncOmni)
-        omni.output_modalities = [None, "audio"]
-        stage = MagicMock()
-        stage.is_comprehension = False
-        omni.stage_list = [stage]
+        omni.engine = SimpleNamespace(supported_tasks=("speech",))
         tasks = await omni.get_supported_tasks()
         assert "generate" not in tasks
         assert "speech" in tasks
@@ -1050,15 +1047,12 @@ class TestAsyncOmniSupportedTasks:
     @pytest.mark.asyncio
     async def test_omni_model_includes_generate(self):
         """Models with text output (e.g. Qwen3-Omni) should include 'generate'."""
-        from unittest.mock import MagicMock
+        from types import SimpleNamespace
 
         from vllm_omni.entrypoints.async_omni import AsyncOmni
 
         omni = AsyncOmni.__new__(AsyncOmni)
-        omni.output_modalities = ["text", None, "audio"]
-        stage = MagicMock()
-        stage.is_comprehension = True
-        omni.stage_list = [stage]
+        omni.engine = SimpleNamespace(supported_tasks=("generate", "speech"))
         tasks = await omni.get_supported_tasks()
         assert "generate" in tasks
 

tests/entrypoints/openai_api/test_video_server.py

@@ -44,7 +44,7 @@ class MockVideoResult:
 
 class FakeAsyncOmni:
     def __init__(self):
-        self.stage_list = ["diffusion"]
+        self.stage_configs = [SimpleNamespace(stage_type="diffusion")]
         self.captured_prompt = None
         self.captured_sampling_params_list = None
 

tests/entrypoints/test_async_omni.py

@@ -0,0 +1,65 @@
+from types import SimpleNamespace
+
+import pytest
+from vllm.entrypoints.openai.models.protocol import BaseModelPath
+from vllm.entrypoints.openai.models.serving import OpenAIServingModels
+
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.asyncio
+async def test_get_supported_tasks_returns_engine_supported_tasks():
+    omni = object.__new__(AsyncOmni)
+    omni.engine = SimpleNamespace(supported_tasks=("generate", "speech"))
+
+    supported_tasks = await omni.get_supported_tasks()
+
+    assert supported_tasks == ("generate", "speech")
+
+
+def test_model_config_and_vllm_config_forward_from_comprehension_stage():
+    model_config = SimpleNamespace(model="Qwen/Qwen3-TTS")
+    vllm_config = SimpleNamespace(model_config=model_config)
+    renderer = SimpleNamespace(name="renderer")
+    input_processor = SimpleNamespace(renderer=renderer)
+    io_processor = SimpleNamespace(name="io-processor")
+    omni = object.__new__(AsyncOmni)
+    omni.engine = SimpleNamespace(
+        stage_clients=[SimpleNamespace(is_comprehension=False), SimpleNamespace(is_comprehension=True)],
+        stage_vllm_configs=[None, vllm_config],
+    )
+    omni.input_processor = input_processor
+    omni.io_processor = io_processor
+
+    assert omni.vllm_config is vllm_config
+    assert omni.model_config is model_config
+    assert omni.renderer is renderer
+    assert omni.input_processor is input_processor
+    assert omni.io_processor is io_processor
+
+
+def test_openai_serving_models_can_consume_async_omni_compat_attrs():
+    model_config = SimpleNamespace(model="Qwen/Qwen3-TTS", max_model_len=32768)
+    vllm_config = SimpleNamespace(model_config=model_config)
+    renderer = SimpleNamespace(name="renderer")
+    input_processor = SimpleNamespace(renderer=renderer)
+    io_processor = SimpleNamespace(name="io-processor")
+    omni = object.__new__(AsyncOmni)
+    omni.engine = SimpleNamespace(
+        stage_clients=[SimpleNamespace(is_comprehension=True)],
+        stage_vllm_configs=[vllm_config],
+    )
+    omni.input_processor = input_processor
+    omni.io_processor = io_processor
+
+    serving_models = OpenAIServingModels(
+        engine_client=omni,
+        base_model_paths=[BaseModelPath(name="tts-model", model_path="Qwen/Qwen3-TTS")],
+    )
+
+    assert serving_models.model_config is model_config
+    assert serving_models.renderer is renderer
+    assert serving_models.io_processor is io_processor
+    assert serving_models.input_processor is input_processor

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -3,87 +3,51 @@
 
 import pytest
 
-from vllm_omni.entrypoints import utils as utils_module
-from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
-MODEL = "riverclouds/qwen_image_random"
+
+def test_default_stage_config_includes_cache_backend():
+    """Ensure cache knobs survive the default diffusion-stage builder."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "cache_backend": "cache_dit",
+            "cache_config": '{"Fn_compute_blocks": 2}',
+            "vae_use_slicing": True,
+            "ulysses_degree": 2,
+        }
+    )[0]
+
+    engine_args = stage_cfg["engine_args"]
+    assert stage_cfg["stage_type"] == "diffusion"
+    assert engine_args["cache_backend"] == "cache_dit"
+    assert engine_args["cache_config"]["Fn_compute_blocks"] == 2
+    assert engine_args["vae_use_slicing"] is True
+    assert engine_args["parallel_config"].ulysses_degree == 2
+    assert engine_args["model_stage"] == "diffusion"
 
 
-def _noop_inline_engine(self, model, stage_config, kwargs):
-    self._inline_diffusion = False
-    self._inline_engine = None
+def test_default_cache_config_used_when_missing():
+    """Ensure default cache_config is synthesized when only backend is given."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "cache_backend": "cache_dit",
+        }
+    )[0]
 
-
-def test_default_stage_config_includes_cache_backend(monkeypatch):
-    """Ensure cache_backend/cache_config are preserved in default diffusion stage."""
-    monkeypatch.setattr(utils_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
-    monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
-    monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
-    monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
-    monkeypatch.setattr(AsyncOmni, "_init_inline_diffusion_engine", _noop_inline_engine)
-
-    omni = AsyncOmni(
-        model=MODEL,
-        cache_backend="cache_dit",
-        cache_config='{"Fn_compute_blocks": 2}',
-        vae_use_slicing=True,
-        ulysses_degree=2,
-    )
-
-    stage_cfg = omni.stage_configs[0]
-    engine_args = stage_cfg.engine_args
-
-    assert engine_args.get("cache_backend") == "cache_dit"
-    cache_config = engine_args.get("cache_config")
-    assert cache_config["Fn_compute_blocks"] == 2
-    assert engine_args.get("vae_use_slicing") is True
-    parallel_config = engine_args.get("parallel_config")
-    if hasattr(parallel_config, "get"):
-        ulysses_degree = parallel_config.get("ulysses_degree")
-    else:
-        ulysses_degree = getattr(parallel_config, "ulysses_degree", None)
-    assert ulysses_degree == 2
-
-
-def test_default_cache_config_used_when_missing(monkeypatch):
-    """Ensure default cache_config is applied when cache_backend is set."""
-    monkeypatch.setattr(utils_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
-    monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
-    monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
-    monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
-    monkeypatch.setattr(AsyncOmni, "_init_inline_diffusion_engine", _noop_inline_engine)
-
-    omni = AsyncOmni(
-        model=MODEL,
-        cache_backend="cache_dit",
-    )
-
-    engine_args = omni.stage_configs[0].engine_args
-    cache_config = engine_args.get("cache_config")
+    cache_config = stage_cfg["engine_args"]["cache_config"]
     assert cache_config is not None
     assert cache_config["Fn_compute_blocks"] == 1
 
 
-def test_default_stage_devices_from_sequence_parallel(monkeypatch):
-    """Ensure devices list reflects sequence parallel size when no parallel_config is provided."""
-    monkeypatch.setattr(utils_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
-    monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
-    monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
-    monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
-    monkeypatch.setattr(AsyncOmni, "_init_inline_diffusion_engine", _noop_inline_engine)
+def test_default_stage_devices_from_sequence_parallel():
+    """Ensure runtime devices reflect computed diffusion world size."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "ulysses_degree": 2,
+            "ring_degree": 2,
+        }
+    )[0]
 
-    omni = AsyncOmni(
-        model=MODEL,
-        ulysses_degree=2,
-        ring_degree=2,
-    )
-
-    stage_cfg = omni.stage_configs[0]
-    runtime = stage_cfg.runtime
-    if hasattr(runtime, "get"):
-        devices = runtime.get("devices")
-    else:
-        devices = getattr(runtime, "devices", None)
-    assert devices == "0,1,2,3"
+    assert stage_cfg["runtime"]["devices"] == "0,1,2,3"

tests/entrypoints/test_omni_input_preprocessor.py

@@ -1,63 +0,0 @@
-import pytest
-
-from vllm_omni.inputs.preprocess import OmniInputPreprocessor
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _make_preprocessor(monkeypatch):
-    preprocessor = object.__new__(OmniInputPreprocessor)
-    monkeypatch.setattr(preprocessor, "_truncate_inputs", lambda tokens, tokenization_kwargs=None: tokens)
-    monkeypatch.setattr(
-        preprocessor,
-        "_process_multimodal",
-        lambda *args, **kwargs: {"prompt_token_ids": [1, 2, 3]},
-    )
-    monkeypatch.setattr(preprocessor, "_tokenize_prompt", lambda prompt_text, tokenization_kwargs=None: [9, 8, 7])
-    return preprocessor
-
-
-def test_process_tokens_keeps_additional_information(monkeypatch):
-    preprocessor = _make_preprocessor(monkeypatch)
-    parsed = {
-        "prompt_token_ids": [1, 2, 3],
-        "prompt_embeds": "embeds",
-        "additional_information": {"task": ["tts"], "lang": ["auto"]},
-    }
-
-    inputs = OmniInputPreprocessor._process_tokens(preprocessor, parsed)
-
-    assert inputs["prompt_token_ids"] == [1, 2, 3]
-    assert inputs["prompt_embeds"] == "embeds"
-    assert inputs["additional_information"] == {"task": ["tts"], "lang": ["auto"]}
-
-
-def test_process_text_keeps_additional_information(monkeypatch):
-    preprocessor = _make_preprocessor(monkeypatch)
-    parsed = {
-        "prompt": "hello",
-        "prompt_embeds": "embeds",
-        "additional_information": {"speaker": ["alice"]},
-    }
-
-    inputs = OmniInputPreprocessor._process_text(preprocessor, parsed)
-
-    assert inputs["prompt_token_ids"] == [9, 8, 7]
-    assert inputs["prompt_embeds"] == "embeds"
-    assert inputs["additional_information"] == {"speaker": ["alice"]}
-
-
-def test_process_text_multimodal_skips_empty_payloads(monkeypatch):
-    preprocessor = _make_preprocessor(monkeypatch)
-    parsed = {
-        "prompt": "hello",
-        "multi_modal_data": {"image": "fake"},
-        "prompt_embeds": None,
-        "additional_information": None,
-    }
-
-    inputs = OmniInputPreprocessor._process_text(preprocessor, parsed)
-
-    assert inputs["prompt_token_ids"] == [1, 2, 3]
-    assert "prompt_embeds" not in inputs
-    assert "additional_information" not in inputs

tests/entrypoints/test_omni_stage_diffusion_config.py

@@ -1,31 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-
-from vllm_omni.entrypoints.omni_stage import _build_od_config
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_build_od_config_includes_diffusion_fields():
-    engine_args = {
-        "cache_backend": "cache_dit",
-        "cache_config": {"Fn_compute_blocks": 2},
-        "vae_use_slicing": True,
-    }
-    od_config = _build_od_config(engine_args, model="dummy-model")
-
-    assert od_config["model"] == "dummy-model"
-    assert od_config["cache_backend"] == "cache_dit"
-    assert od_config["cache_config"]["Fn_compute_blocks"] == 2
-    assert od_config["vae_use_slicing"] is True
-
-
-def test_build_od_config_respects_explicit_config():
-    engine_args = {
-        "od_config": {"cache_backend": "tea_cache"},
-        "cache_backend": "cache_dit",
-    }
-    od_config = _build_od_config(engine_args, model="dummy-model")
-    assert od_config == {"cache_backend": "tea_cache"}

tests/perf/scripts/run_benchmark.py

@@ -50,11 +50,14 @@ def create_unique_server_params(configs: list[dict[str, Any]]) -> list[tuple[str
     for config in configs:
         test_name = config["test_name"]
         model = config["server_params"]["model"]
-        stage_config_name = config["server_params"]["stage_config_name"]
-        stage_config_path = str(Path(__file__).parent.parent / "stage_configs" / stage_config_name)
-        delete = config["server_params"].get("delete", None)
-        update = config["server_params"].get("update", None)
-        stage_config_path = modify_stage(stage_config_path, update, delete)
+        stage_config_name = config["server_params"].get("stage_config_name")
+        if stage_config_name:
+            stage_config_path = str(Path(__file__).parent.parent / "stage_configs" / stage_config_name)
+            delete = config["server_params"].get("delete", None)
+            update = config["server_params"].get("update", None)
+            stage_config_path = modify_stage(stage_config_path, update, delete)
+        else:
+            stage_config_path = None
 
         server_param = (test_name, model, stage_config_path)
         if server_param not in seen:
@@ -98,7 +101,10 @@ def omni_server(request):
 
         print(f"Starting OmniServer with test: {test_name}, model: {model}")
 
-        with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server:
+        server_args = ["--stage-init-timeout", "120"]
+        if stage_config_path:
+            server_args = ["--stage-configs-path", stage_config_path] + server_args
+        with OmniServer(model, server_args) as server:
             server.test_name = test_name
             print("OmniServer started successfully")
             yield server

tests/perf/stage_configs/qwen3_omni.yaml

@@ -7,7 +7,7 @@
 async_chunk: false
 stage_args:
   - stage_id: 0
-    stage_type: llm  # Use llm stage type to launch OmniLLM
+    stage_type: llm  # Use llm stage type for AR stages
     runtime:
       devices: "0"
       max_batch_size: 64
@@ -38,7 +38,7 @@ stage_args:
       repetition_penalty: 1.05
 
   - stage_id: 1
-    stage_type: llm  # Use llm stage type to launch OmniLLM
+    stage_type: llm  # Use llm stage type for AR stages
     runtime:
       devices: "1"
       max_batch_size: 64
@@ -69,7 +69,7 @@ stage_args:
       stop_token_ids: [2150]
 
   - stage_id: 2
-    stage_type: llm  # Use llm stage type to launch OmniLLM
+    stage_type: llm  # Use llm stage type for AR stages
     runtime:
       devices: "1"
       max_batch_size: 64

tests/perf/stage_configs/qwen3_tts.yaml

@@ -7,6 +7,7 @@ async_chunk: true
 stage_args:
   - stage_id: 0
     stage_type: llm
+    is_comprehension: true
     runtime:
       devices: "0"
       max_batch_size: 4

tests/perf/tests/test.json

@@ -152,8 +152,7 @@
     {
         "test_name": "test_qwen3_tts",
         "server_params": {
-            "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
-            "stage_config_name": "qwen3_tts.yaml"
+            "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
         },
         "benchmark_params": [
             {

tools/pre_commit/check_pickle_imports.py

@@ -16,7 +16,6 @@ import regex as re
 #  alternatives like msgpack or pydantic that are already in use in vLLM. Only
 #  add to this list if absolutely necessary and after careful security review.
 ALLOWED_FILES = {
-    "vllm_omni/entrypoints/omni_llm.py",
     "tests/e2e/offline_inference/utils.py",
     "tests/utils.py",
     "vllm_omni/diffusion/distributed/group_coordinator.py",

vllm_omni/__init__.py

@@ -24,10 +24,7 @@ except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
 from vllm_omni.transformers_utils import configs as _configs  # noqa: F401, E402
 
 from .config import OmniModelConfig
-from .entrypoints.async_omni import AsyncOmni
-
-# Main entry points
-from .entrypoints.omni import Omni
+from .entrypoints import AsyncOmni, Omni
 
 from .version import __version__, __version_tuple__  # isort:skip
 

vllm_omni/config/stage_config.py

@@ -49,10 +49,11 @@ class StageType(str, Enum):
 
 @dataclass
 class StageConfig:
-    """Per-stage configuration — pipeline-structure fields only.
+    """Per-stage configuration from pipeline YAML.
 
-    Engine params (gpu_memory_utilization, tp_size, etc.) come from CLI,
-    NOT from this class.
+    Topology fields (stage_id, input_sources, etc.) define the DAG.
+    Engine and runtime defaults come from the YAML; CLI overrides take
+    precedence via ``runtime_overrides``.
     """
 
     # Identity
@@ -71,6 +72,14 @@ class StageConfig:
     hf_config_name: str | None = None
     is_comprehension: bool = False
 
+    # Per-stage engine args from pipeline YAML (defaults)
+    yaml_engine_args: dict[str, Any] = field(default_factory=dict)
+    # Per-stage runtime config from pipeline YAML (devices, max_batch_size)
+    yaml_runtime: dict[str, Any] = field(default_factory=dict)
+    # Pass-through fields from pipeline YAML (default_sampling_params,
+    # output_connectors, input_connectors, tts_args, etc.)
+    yaml_extras: dict[str, Any] = field(default_factory=dict)
+
     # Runtime overrides (populated from CLI, not from pipeline YAML)
     runtime_overrides: dict[str, Any] = field(default_factory=dict)
 
@@ -80,11 +89,11 @@ class StageConfig:
         Returns:
             OmegaConf DictConfig with stage configuration in legacy format.
         """
-        # Build engine_args dict with required fields
-        engine_args: dict[str, Any] = {
-            "model_stage": self.model_stage,
-        }
+        # Start with YAML engine_args defaults
+        engine_args: dict[str, Any] = dict(self.yaml_engine_args)
 
+        # Overlay topology-level fields
+        engine_args["model_stage"] = self.model_stage
         if self.worker_type:
             engine_args["worker_type"] = self.worker_type
         if self.scheduler_cls:
@@ -92,19 +101,22 @@ class StageConfig:
         if self.hf_config_name:
             engine_args["hf_config_name"] = self.hf_config_name
 
-        # Apply runtime overrides (CLI args)
+        # CLI overrides take precedence over YAML defaults
         for key, value in self.runtime_overrides.items():
             if key not in ("devices", "max_batch_size"):
                 engine_args[key] = value
 
-        # Build runtime config
-        runtime: dict[str, Any] = {
-            "process": True,
-            "max_batch_size": self.runtime_overrides.get("max_batch_size", 1),
-        }
+        # Build runtime config from YAML defaults + CLI overrides
+        runtime: dict[str, Any] = dict(self.yaml_runtime)
+        runtime.setdefault("process", True)
+        runtime.setdefault("max_batch_size", self.runtime_overrides.get("max_batch_size", 1))
         if "devices" in self.runtime_overrides:
             runtime["devices"] = self.runtime_overrides["devices"]
 
+        # Inject max_num_seqs from max_batch_size (legacy compat)
+        max_batch_size = int(runtime.get("max_batch_size", 1) or 1)
+        engine_args.setdefault("max_num_seqs", max_batch_size)
+
         # Build full config dict
         config_dict: dict[str, Any] = {
             "stage_id": self.stage_id,
@@ -120,6 +132,10 @@ class StageConfig:
         if self.custom_process_input_func:
             config_dict["custom_process_input_func"] = self.custom_process_input_func
 
+        # Pass through extra YAML fields (default_sampling_params,
+        # output_connectors, input_connectors, tts_args, etc.)
+        config_dict.update(self.yaml_extras)
+
         return create_config(config_dict)
 
 
@@ -249,6 +265,15 @@ class StageConfigFactory:
         if errors:
             logger.warning(f"Pipeline validation warnings for {model}: {errors}")
 
+        # Inject pipeline-wide async_chunk into ALL stages' engine_args.
+        # The legacy loader (load_stage_configs_from_yaml) sets async_chunk
+        # on every stage so that build_engine_args_dict() can inject the
+        # stage_connector_spec.  AsyncOmniEngine.__init__ also reads it
+        # from stage_configs[0].engine_args.async_chunk.
+        if pipeline.async_chunk:
+            for stage in pipeline.stages:
+                stage.yaml_engine_args.setdefault("async_chunk", True)
+
         # Apply CLI overrides
         result: list[StageConfig] = []
         for stage in pipeline.stages:
@@ -348,6 +373,25 @@ class StageConfigFactory:
 
         return cls._parse_pipeline_yaml(pipeline_path, model_type)
 
+    # Keys consumed as explicit StageConfig fields — everything else is
+    # passed through via yaml_extras.
+    _KNOWN_STAGE_KEYS: set[str] = {
+        "stage_id",
+        "model_stage",
+        "stage_type",
+        "input_sources",
+        "engine_input_source",
+        "custom_process_input_func",
+        "final_output",
+        "final_output_type",
+        "worker_type",
+        "scheduler_cls",
+        "hf_config_name",
+        "is_comprehension",
+        "engine_args",
+        "runtime",
+    }
+
     @classmethod
     def _parse_pipeline_yaml(cls, path: Path, model_type: str) -> ModelPipeline:
         """Parse a pipeline YAML file.
@@ -375,27 +419,65 @@ class StageConfigFactory:
                 input_sources = []
             input_sources = list(input_sources)
 
+            # Extract per-stage engine_args and runtime dicts
+            raw_ea = stage_data.get("engine_args", None)
+            yaml_engine_args = to_dict(raw_ea) if raw_ea is not None else {}
+            raw_rt = stage_data.get("runtime", None)
+            yaml_runtime = to_dict(raw_rt) if raw_rt is not None else {}
+
+            # Topology-level fields that also live inside engine_args in legacy
+            # YAMLs (worker_type, scheduler_cls, etc.) — read from both places.
+            worker_type = stage_data.get("worker_type", None) or yaml_engine_args.pop("worker_type", None)
+            scheduler_cls = stage_data.get("scheduler_cls", None) or yaml_engine_args.pop("scheduler_cls", None)
+            hf_config_name = stage_data.get("hf_config_name", None) or yaml_engine_args.pop("hf_config_name", None)
+            model_stage = getattr(stage_data, "model_stage", None) or yaml_engine_args.pop("model_stage", None)
+
+            # Collect pass-through fields (default_sampling_params,
+            # output_connectors, input_connectors, tts_args, etc.)
+            yaml_extras: dict[str, Any] = {}
+            for key in stage_data:
+                if key not in cls._KNOWN_STAGE_KEYS:
+                    val = stage_data[key]
+                    try:
+                        yaml_extras[key] = to_dict(val)
+                    except ValueError:
+                        yaml_extras[key] = val
+
             stage = StageConfig(
                 stage_id=stage_data.stage_id,
-                model_stage=stage_data.model_stage,
+                model_stage=model_stage or "",
                 stage_type=stage_type,
                 input_sources=input_sources,
                 custom_process_input_func=stage_data.get("custom_process_input_func", None),
                 final_output=stage_data.get("final_output", False),
                 final_output_type=stage_data.get("final_output_type", None),
-                worker_type=stage_data.get("worker_type", None),
-                scheduler_cls=stage_data.get("scheduler_cls", None),
-                hf_config_name=stage_data.get("hf_config_name", None),
+                worker_type=worker_type,
+                scheduler_cls=scheduler_cls,
+                hf_config_name=hf_config_name,
                 is_comprehension=stage_data.get("is_comprehension", False),
+                yaml_engine_args=yaml_engine_args,
+                yaml_runtime=yaml_runtime,
+                yaml_extras=yaml_extras,
             )
             stages.append(stage)
 
         # Get pipeline-wide flags
         async_chunk = config_data.get("async_chunk", False)
 
-        # Get optional connector config
-        connectors = to_dict(config_data.connectors) if hasattr(config_data, "connectors") else None
-        edges = to_dict(config_data.edges) if hasattr(config_data, "edges") else None
+        # Get optional connector config — check both top-level and nested
+        # under ``runtime`` (legacy stage_configs format).
+        connectors = None
+        edges = None
+        if hasattr(config_data, "connectors"):
+            connectors = to_dict(config_data.connectors)
+        if hasattr(config_data, "edges"):
+            edges = to_dict(config_data.edges)
+        if hasattr(config_data, "runtime") and config_data.runtime is not None:
+            top_runtime = config_data.runtime
+            if connectors is None and hasattr(top_runtime, "connectors"):
+                connectors = to_dict(top_runtime.connectors)
+            if edges is None and hasattr(top_runtime, "edges"):
+                edges = to_dict(top_runtime.edges)
 
         return ModelPipeline(
             model_type=getattr(config_data, "model_type", model_type),
@@ -421,9 +503,21 @@ class StageConfigFactory:
 
             hf_config = get_config(model, trust_remote_code=trust_remote_code)
             return hf_config.model_type, hf_config
+        except Exception:
+            pass
+
+        # Fallback: read config.json directly for custom model types that
+        # are not registered with transformers (e.g. qwen3_tts).
+        try:
+            from vllm.transformers_utils.config import get_hf_file_to_dict
+
+            config_dict = get_hf_file_to_dict("config.json", model, revision=None)
+            if config_dict and "model_type" in config_dict:
+                return config_dict["model_type"], None
         except Exception as e:
             logger.debug(f"Failed to auto-detect model type for {model}: {e}")
-            return None, None
+
+        return None, None
 
     # Keys that should never be forwarded as engine overrides (internal /
     # orchestrator-only knobs, complex objects, etc.).

vllm_omni/diffusion/stage_diffusion_client.py

@@ -0,0 +1,146 @@
+"""Stage Diffusion Client for vLLM-Omni multi-stage runtime.
+
+Wraps AsyncOmniDiffusion to expose the same interface the Orchestrator
+expects from any stage client.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Any
+
+from vllm.logger import init_logger
+
+from vllm_omni.engine.stage_init_utils import StageMetadata
+from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
+from vllm_omni.outputs import OmniRequestOutput
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
+
+logger = init_logger(__name__)
+
+
+class StageDiffusionClient:
+    """Wraps AsyncOmniDiffusion for use inside the Orchestrator.
+
+    Exposes the same attributes and async methods the Orchestrator
+    uses on StageEngineCoreClient, but routes execution through
+    DiffusionEngine instead of vLLM EngineCore.
+    """
+
+    stage_type: str = "diffusion"
+
+    def __init__(
+        self,
+        model: str,
+        od_config: OmniDiffusionConfig,
+        metadata: StageMetadata,
+    ) -> None:
+        self.stage_id = metadata.stage_id
+        self.final_output = metadata.final_output
+        self.final_output_type = metadata.final_output_type
+        self.default_sampling_params = metadata.default_sampling_params
+        self.custom_process_input_func = metadata.custom_process_input_func
+        self.engine_input_source = metadata.engine_input_source
+
+        self._engine = AsyncOmniDiffusion(model=model, od_config=od_config)
+        self._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue()
+        self._tasks: dict[str, asyncio.Task] = {}
+
+        logger.info("[StageDiffusionClient] Stage-%s initialized", self.stage_id)
+
+    async def add_request_async(
+        self,
+        request_id: str,
+        prompt: OmniPromptType,
+        sampling_params: OmniDiffusionSamplingParams,
+    ) -> None:
+        task = asyncio.create_task(
+            self._run(request_id, prompt, sampling_params),
+            name=f"diffusion-{request_id}",
+        )
+        self._tasks[request_id] = task
+
+    async def _run(
+        self,
+        request_id: str,
+        prompt: OmniPromptType,
+        sampling_params: OmniDiffusionSamplingParams,
+    ) -> None:
+        try:
+            result = await self._engine.generate(prompt, sampling_params, request_id)
+            await self._output_queue.put(result)
+        except Exception as e:
+            logger.exception(
+                "[StageDiffusionClient] Stage-%s req=%s failed: %s",
+                self.stage_id,
+                request_id,
+                e,
+            )
+        finally:
+            self._tasks.pop(request_id, None)
+
+    def get_diffusion_output_async(self) -> OmniRequestOutput | None:
+        try:
+            return self._output_queue.get_nowait()
+        except asyncio.QueueEmpty:
+            return None
+
+    async def abort_requests_async(self, request_ids: list[str]) -> None:
+        for rid in request_ids:
+            task = self._tasks.pop(rid, None)
+            if task:
+                task.cancel()
+
+    async def collective_rpc_async(
+        self,
+        method: str,
+        timeout: float | None = None,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> Any:
+        """Best-effort control RPC shim for diffusion stages.
+
+        TODO(AsyncOmni): add dedicated wrappers on AsyncOmniDiffusion for the
+        remaining control APIs instead of reaching into its underlying engine.
+        """
+        kwargs = kwargs or {}
+
+        if method in {"add_lora", "remove_lora", "list_loras", "pin_lora", "start_profile", "stop_profile"}:
+            target = getattr(self._engine, method, None)
+            if target is None:
+                return {
+                    "supported": False,
+                    "todo": True,
+                    "reason": f"AsyncOmniDiffusion.{method} is not implemented",
+                }
+            result = target(*args, **kwargs)
+            if timeout is not None:
+                return await asyncio.wait_for(result, timeout=timeout)
+            return await result
+
+        if method in {"sleep", "wake_up"}:
+            loop = asyncio.get_running_loop()
+            return await loop.run_in_executor(
+                self._engine._executor,
+                self._engine.engine.collective_rpc,
+                method,
+                timeout,
+                args,
+                kwargs,
+                None,
+            )
+
+        return {
+            "supported": False,
+            "todo": True,
+            "reason": f"Diffusion stage collective_rpc method {method} is not implemented yet",
+        }
+
+    def shutdown(self) -> None:
+        for task in self._tasks.values():
+            task.cancel()
+        self._tasks.clear()
+        self._engine.close()

vllm_omni/distributed/omni_connectors/adapter.py

@@ -1,13 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# temporary for compatibility with vllm_omni.entrypoints.omni_stage.py
-# and vllm_omni.entrypoints.omni_llm.py
-
 import time
 from collections.abc import Callable
 from typing import Any
 
-from vllm_omni.entrypoints.stage_utils import OmniStageTaskType
 from vllm_omni.metrics import OrchestratorAggregator
 
 from .utils.logging import get_connector_logger
@@ -66,7 +62,7 @@ def try_send_via_connector(
         if success:
             # Send lightweight notification via queue
             notify_payload = {
-                "type": OmniStageTaskType.GENERATE,
+                "type": "generate",
                 "request_id": req_id,
                 "sampling_params": sampling_params,
                 "from_connector": True,

vllm_omni/engine/arg_utils.py

@@ -1,8 +1,10 @@
+import argparse
+import dataclasses
 from dataclasses import dataclass, field
 from typing import Any
 
 import vllm.envs as envs
-from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
+from vllm.engine.arg_utils import EngineArgs
 from vllm.logger import init_logger
 from vllm.transformers_utils.gguf_utils import is_gguf
 
@@ -86,6 +88,12 @@ class OmniEngineArgs(EngineArgs):
         load_omni_general_plugins()
         super().__post_init__()
 
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace) -> "OmniEngineArgs":
+        attrs = [attr.name for attr in dataclasses.fields(cls)]
+        engine_args = cls(**{attr: getattr(args, attr) for attr in attrs if hasattr(args, attr)})
+        return engine_args
+
     def _ensure_omni_models_registered(self):
         if hasattr(self, "_omni_models_registered"):
             return True
@@ -115,170 +123,8 @@ class OmniEngineArgs(EngineArgs):
         self._ensure_omni_models_registered()
 
         # Keep compatibility when async args are constructed from partial payloads.
-        limit_mm_per_prompt = getattr(self, "limit_mm_per_prompt", {})
-        language_model_only = getattr(self, "language_model_only", False)
-        enable_mm_embeds = getattr(self, "enable_mm_embeds", False)
-        interleave_mm_strings = getattr(self, "interleave_mm_strings", False)
-        media_io_kwargs = getattr(self, "media_io_kwargs", {})
-        skip_mm_profiling = getattr(self, "skip_mm_profiling", False)
-        mm_processor_kwargs = getattr(self, "mm_processor_kwargs", None)
-        mm_processor_cache_gb = getattr(self, "mm_processor_cache_gb", 4)
-        mm_processor_cache_type = getattr(self, "mm_processor_cache_type", None)
-        mm_shm_cache_max_object_size_mb = getattr(self, "mm_shm_cache_max_object_size_mb", 128)
-        mm_encoder_only = getattr(self, "mm_encoder_only", False)
-        mm_encoder_tp_mode = getattr(self, "mm_encoder_tp_mode", "weights")
-        mm_encoder_attn_backend = getattr(self, "mm_encoder_attn_backend", None)
-        video_pruning_rate = getattr(self, "video_pruning_rate", 0.0)
-
-        # Build stage_connector_config from stage_connector_spec
-        stage_connector_config = {
-            "name": self.stage_connector_spec.get("name", "SharedMemoryConnector"),
-            "extra": self.stage_connector_spec.get("extra", {}).copy(),
-        }
-        stage_connector_config["extra"]["stage_id"] = self.stage_id
-
-        # Create OmniModelConfig directly from engine args
-        # Note: We pass the actual init parameters matching vLLM's EngineArgs.create_model_config()
-        omni_config = OmniModelConfig(
-            # Base ModelConfig fields (matching vLLM's EngineArgs.create_model_config)
-            model=self.model,
-            model_weights=self.model_weights,
-            hf_config_path=self.hf_config_path,
-            runner=self.runner,
-            convert=self.convert,
-            tokenizer=self.tokenizer,
-            tokenizer_mode=self.tokenizer_mode,
-            trust_remote_code=self.trust_remote_code,
-            allowed_local_media_path=self.allowed_local_media_path,
-            allowed_media_domains=self.allowed_media_domains,
-            dtype=self.dtype,
-            seed=self.seed,
-            revision=self.revision,
-            code_revision=self.code_revision,
-            hf_token=self.hf_token,
-            hf_overrides=self.hf_overrides,
-            tokenizer_revision=self.tokenizer_revision,
-            max_model_len=self.max_model_len,
-            quantization=self.quantization,
-            allow_deprecated_quantization=self.allow_deprecated_quantization,
-            enforce_eager=self.enforce_eager,
-            enable_return_routed_experts=self.enable_return_routed_experts,
-            max_logprobs=self.max_logprobs,
-            logprobs_mode=self.logprobs_mode,
-            disable_sliding_window=self.disable_sliding_window,
-            disable_cascade_attn=self.disable_cascade_attn,
-            skip_tokenizer_init=self.skip_tokenizer_init,
-            enable_prompt_embeds=self.enable_prompt_embeds,
-            served_model_name=self.served_model_name,
-            language_model_only=language_model_only,
-            limit_mm_per_prompt=limit_mm_per_prompt,
-            enable_mm_embeds=enable_mm_embeds,
-            interleave_mm_strings=interleave_mm_strings,
-            media_io_kwargs=media_io_kwargs,
-            skip_mm_profiling=skip_mm_profiling,
-            config_format=self.config_format,
-            mm_processor_kwargs=mm_processor_kwargs,
-            mm_processor_cache_gb=mm_processor_cache_gb,
-            mm_processor_cache_type=mm_processor_cache_type,
-            mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb,
-            mm_encoder_only=mm_encoder_only,
-            mm_encoder_tp_mode=mm_encoder_tp_mode,
-            mm_encoder_attn_backend=mm_encoder_attn_backend,
-            pooler_config=self.pooler_config,
-            generation_config=self.generation_config,
-            override_generation_config=self.override_generation_config,
-            enable_sleep_mode=self.enable_sleep_mode,
-            model_impl=self.model_impl,
-            override_attention_dtype=self.override_attention_dtype,
-            logits_processors=self.logits_processors,
-            video_pruning_rate=video_pruning_rate,
-            io_processor_plugin=self.io_processor_plugin,
-            # Omni-specific fields
-            stage_id=self.stage_id,
-            async_chunk=self.async_chunk,
-            model_stage=self.model_stage,
-            model_arch=self.model_arch,
-            worker_type=self.worker_type,
-            engine_output_type=self.engine_output_type,
-            hf_config_name=self.hf_config_name,
-            custom_process_next_stage_input_func=self.custom_process_next_stage_input_func,
-            stage_connector_config=stage_connector_config,
-            omni_kv_config=self.omni_kv_config,
-            task_type=self.task_type,
-        )
-        omni_config.hf_config.architectures = omni_config.architectures
-
-        return omni_config
-
-
-@dataclass
-class AsyncOmniEngineArgs(AsyncEngineArgs):
-    """Async engine arguments for omni models, extending base AsyncEngineArgs.
-    Adds omni-specific configuration fields for multi-stage pipeline
-    processing and output type specification in async contexts.
-    Args:
-        stage_id: Identifier for the stage in a multi-stage pipeline (default: 0)
-        model_stage: Stage type identifier, e.g., "thinker" or "talker"
-            (default: "thinker")
-        model_arch: Model architecture name
-            (default: "Qwen2_5OmniForConditionalGeneration")
-        engine_output_type: Optional output type specification for the engine.
-            Used to route outputs to appropriate processors (e.g., "image",
-            "audio", "latents"). If None, output type is inferred.
-        stage_connector_spec: Extra configuration for stage connector
-        worker_type: Model Type, e.g., "ar" or "generation"
-        task_type: Default task type for TTS models (CustomVoice, VoiceDesign, or Base).
-            If not specified, will be inferred from model path.
-    """
-
-    stage_id: int = 0
-    model_stage: str = "thinker"
-    model_arch: str | None = None
-    engine_output_type: str | None = None
-    hf_config_name: str | None = None
-    custom_process_next_stage_input_func: str | None = None
-    stage_connector_spec: dict[str, Any] = field(default_factory=dict)
-    async_chunk: bool = False
-    omni_kv_config: dict | None = None
-    quantization_config: Any | None = None
-    worker_type: str | None = None
-    task_type: str | None = None
-
-    def __post_init__(self) -> None:
-        load_omni_general_plugins()
-        super().__post_init__()
-
-    def _ensure_omni_models_registered(self):
-        if hasattr(self, "_omni_models_registered"):
-            return True
-        register_omni_models_to_vllm()
-        self._omni_models_registered = True
-        return True
-
-    def create_model_config(self) -> OmniModelConfig:
-        """Create an OmniModelConfig from these engine arguments.
-        Returns:
-            OmniModelConfig instance with all configuration fields set
-        """
-        # GGUF files need a specific model loader path in vLLM.
-        if is_gguf(self.model):
-            self.quantization = self.load_format = "gguf"
-
-        if not envs.VLLM_ENABLE_V1_MULTIPROCESSING:
-            logger.warning(
-                "The global random seed is set to %d. Since "
-                "VLLM_ENABLE_V1_MULTIPROCESSING is set to False, this may "
-                "affect the random state of the Python process that "
-                "launched vLLM.",
-                self.seed,
-            )
-
-        # register omni models to avoid model not found error
-        self._ensure_omni_models_registered()
-
-        # Keep compatibility when async args are constructed from partial payloads.
-        limit_mm_per_prompt = getattr(self, "limit_mm_per_prompt", {})
         language_model_only = getattr(self, "language_model_only", False)
+        limit_mm_per_prompt = getattr(self, "limit_mm_per_prompt", {})
         enable_mm_embeds = getattr(self, "enable_mm_embeds", False)
         interleave_mm_strings = getattr(self, "interleave_mm_strings", False)
         media_io_kwargs = getattr(self, "media_io_kwargs", {})

vllm_omni/engine/input_processor.py

@@ -1,300 +0,0 @@
-import time
-from collections.abc import Mapping
-from typing import Any
-
-import numpy as np
-import torch
-from vllm.config import VllmConfig
-from vllm.inputs import ProcessorInputs, PromptType
-from vllm.inputs.parse import split_enc_dec_inputs
-from vllm.logger import init_logger
-from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
-from vllm.multimodal.inputs import MultiModalFeatureSpec
-from vllm.multimodal.utils import argsort_mm_positions
-from vllm.pooling_params import PoolingParams
-from vllm.renderers import BaseRenderer
-from vllm.sampling_params import SamplingParams
-from vllm.tasks import SupportedTask
-from vllm.utils import length_from_prompt_token_ids_or_embeds
-from vllm.utils.jsontree import json_iter_leaves
-from vllm.v1.engine.input_processor import InputProcessor
-
-from vllm_omni.engine import (
-    AdditionalInformationEntry,
-    AdditionalInformationPayload,
-    OmniEngineCoreRequest,
-    PromptEmbedsPayload,
-)
-from vllm_omni.inputs.preprocess import OmniInputPreprocessor
-from vllm_omni.lora.request import LoRARequest
-
-logger = init_logger(__name__)
-
-_OMNI_EXTRA_KEYS = (
-    "additional_information",
-    "prompt_embeds",
-    "negative_prompt",
-    "negative_prompt_embeds",
-)
-
-
-def reinject_omni_fields(
-    results: list[ProcessorInputs],
-    original_prompts: list[dict],
-) -> None:
-    """Re-inject omni-specific fields that the upstream renderer discards.
-
-    The upstream renderer's ``process_for_engine`` creates new dicts that only
-    copy standard vLLM fields (prompt_token_ids, multi_modal_data, …).
-    Omni-specific fields such as ``additional_information`` and
-    ``prompt_embeds`` are silently dropped.  This helper copies them back from
-    the *original* parsed prompts into the renderer outputs so they survive
-    into ``OmniInputProcessor.process_inputs()``.
-    """
-    for result, orig in zip(results, original_prompts):
-        if not isinstance(orig, dict):
-            continue
-        for key in _OMNI_EXTRA_KEYS:
-            val = orig.get(key)
-            if val is not None and key not in result:
-                result[key] = val
-
-
-class OmniInputProcessor(InputProcessor):
-    """Processor for omni models, handling multimodal inputs and embeddings.
-
-    Extends the base vLLM Processor with support for processing prompt
-    embeddings and additional information payloads, enabling direct transfer
-    of pre-computed embeddings between pipeline stages.
-
-    Args:
-        vllm_config: Global vLLM configuration
-        mm_registry: Multi-modal registry for processing multimodal inputs
-    """
-
-    @staticmethod
-    def _dtype_to_name(dtype: torch.dtype) -> str:
-        """Convert torch dtype to string representation.
-
-        Args:
-            dtype: PyTorch dtype to convert
-
-        Returns:
-            String representation of the dtype (e.g., "float32", "int64")
-        """
-        mapping = {
-            torch.float32: "float32",
-            torch.float: "float32",
-            torch.float16: "float16",
-            torch.half: "float16",
-            torch.bfloat16: "bfloat16",
-            torch.float64: "float64",
-            torch.double: "float64",
-            torch.int64: "int64",
-            torch.long: "int64",
-            torch.int32: "int32",
-            torch.int: "int32",
-            torch.int16: "int16",
-            torch.short: "int16",
-            torch.int8: "int8",
-            torch.uint8: "uint8",
-            torch.bool: "bool",
-        }
-        return mapping.get(dtype, str(dtype).replace("torch.", ""))
-
-    def __init__(
-        self,
-        vllm_config: VllmConfig,
-        renderer: BaseRenderer | None = None,
-        *,
-        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
-    ):
-        super().__init__(vllm_config, renderer=renderer, mm_registry=mm_registry)
-        self.input_preprocessor = OmniInputPreprocessor(
-            vllm_config=vllm_config,
-            renderer=self.renderer,
-            mm_registry=mm_registry,
-        )
-
-    def process_inputs(
-        self,
-        request_id: str,
-        prompt: PromptType | ProcessorInputs,
-        params: SamplingParams | PoolingParams,
-        supported_tasks: tuple[SupportedTask, ...] = ("generate",),
-        arrival_time: float | None = None,
-        lora_request: LoRARequest | None = None,
-        tokenization_kwargs: dict[str, Any] | None = None,
-        trace_headers: Mapping[str, str] | None = None,
-        priority: int = 0,
-        data_parallel_rank: int | None = None,
-        resumable: bool = False,
-    ) -> OmniEngineCoreRequest:
-        """Process input prompt into an engine core request.
-
-        Converts a prompt (text, tokens, or multimodal) into an
-        OmniEngineCoreRequest that can be processed by the engine.
-        Handles prompt embeddings and additional information payloads
-        for direct transfer between stages.
-
-        Args:
-            request_id: Unique identifier for this request
-            prompt: Input prompt (text, token IDs, embeddings, or multimodal)
-            params: Sampling or pooling parameters for generation
-            supported_tasks: Tuple of supported tasks for validation
-            arrival_time: Optional arrival timestamp (defaults to current time)
-            lora_request: Optional LoRA adapter request
-            tokenization_kwargs: Optional additional tokenization arguments
-            trace_headers: Optional tracing headers for observability
-            priority: Request priority (higher values processed first)
-            data_parallel_rank: Optional data parallel rank for distributed
-                inference
-            resumable: Whether the request supports streaming input
-
-        Returns:
-            OmniEngineCoreRequest ready for the engine
-
-        Raises:
-            ValueError: If data_parallel_rank is out of range or prompt_embeds
-                has incorrect shape
-        """
-        self._validate_params(params, supported_tasks)
-        self._validate_lora(lora_request)
-
-        parallel_config = self.vllm_config.parallel_config
-        dp_size = parallel_config.data_parallel_size
-        dp_local_size = parallel_config.data_parallel_size_local
-        num_ranks = dp_local_size if parallel_config.local_engines_only else dp_size
-        if data_parallel_rank is not None and not (0 <= data_parallel_rank < num_ranks):
-            raise ValueError(f"data_parallel_rank {data_parallel_rank} is out of range [0, {num_ranks}).")
-
-        # Short-circuit for prompts already processed by the renderer
-        # (they carry a "type" key).  Raw prompts must still go through the
-        # omni preprocessor which preserves additional_information, etc.
-        if isinstance(prompt, dict) and "type" in prompt:
-            if arrival_time is None:
-                arrival_time = prompt.get("arrival_time", time.time())
-            processed_inputs: ProcessorInputs = prompt  # type: ignore[assignment]
-        else:
-            if arrival_time is None:
-                arrival_time = time.time()
-
-            processed_inputs = self.input_preprocessor.preprocess(
-                prompt,
-                tokenization_kwargs=tokenization_kwargs,
-            )
-
-        self._platform_validate_request(processed_inputs, params)
-
-        encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs)
-        self._validate_model_inputs(encoder_inputs, decoder_inputs)
-
-        # Normalize decoder prompt access across TypedDict variants.
-        if decoder_inputs["type"] == "embeds":
-            prompt_token_ids = None
-            prompt_embeds = decoder_inputs["prompt_embeds"]
-        else:
-            prompt_token_ids = decoder_inputs["prompt_token_ids"]
-            prompt_embeds = decoder_inputs.get("prompt_embeds")
-
-        sampling_params = None
-        pooling_params = None
-        if isinstance(params, SamplingParams):
-            # TODO: can we avoid cloning here in multiproc case?
-            sampling_params = params.clone()
-            # If unset max tokens, then generate up to the max_model_len.
-            if sampling_params.max_tokens is None:
-                seq_len = length_from_prompt_token_ids_or_embeds(prompt_token_ids, prompt_embeds)
-                sampling_params.max_tokens = self.model_config.max_model_len - seq_len
-            sampling_params.update_from_generation_config(
-                self.generation_config_fields,
-                self.renderer.get_eos_token_id(),
-            )
-            if self.tokenizer is not None:
-                sampling_params.update_from_tokenizer(self.tokenizer)
-        else:
-            pooling_params = params.clone()
-
-        # Multimodal related.
-        mm_features: list[MultiModalFeatureSpec] | None = None
-
-        if decoder_inputs["type"] == "multimodal":
-            decoder_mm_inputs = decoder_inputs["mm_kwargs"]
-            decoder_mm_positions = decoder_inputs["mm_placeholders"]
-            decoder_mm_hashes = decoder_inputs["mm_hashes"]
-
-            if not all(isinstance(leaf, str) for leaf in json_iter_leaves(decoder_mm_hashes)):
-                raise ValueError(
-                    f"mm_hashes must contain only strings, got: {decoder_mm_hashes}. "
-                    "This is likely due to an incorrect custom implementation of "
-                    "MultiModalProcessor.apply method."
-                )
-
-            # Merge and flatten multimodal placeholders, hashes and inputs
-            # from dictionaries to lists, and sort them by each item's position
-            # in the input sequence.
-            sorted_mm_idxs = argsort_mm_positions(decoder_mm_positions)
-
-            mm_features = []
-            for modality, idx in sorted_mm_idxs:
-                base_mm_hash = decoder_mm_hashes[modality][idx]
-                mm_features.append(
-                    MultiModalFeatureSpec(
-                        data=decoder_mm_inputs[modality][idx],
-                        modality=modality,
-                        identifier=self._get_mm_identifier(base_mm_hash, lora_request),
-                        mm_position=decoder_mm_positions[modality][idx],
-                        mm_hash=base_mm_hash,
-                    )
-                )
-
-        # Compatibility: decode serialized prompt embeds if provided.
-        if isinstance(prompt_embeds, PromptEmbedsPayload):
-            prompt_embeds = self._decode_prompt_embeds(prompt_embeds)
-
-        additional_information_payload: AdditionalInformationPayload | None = None
-        raw_info: dict[str, Any] | AdditionalInformationPayload | None = decoder_inputs.get("additional_information")
-        if isinstance(raw_info, AdditionalInformationPayload):
-            additional_information_payload = raw_info
-        elif raw_info is not None:
-            entries: dict[str, AdditionalInformationEntry] = {}
-            for key, value in raw_info.items():
-                if isinstance(value, torch.Tensor):
-                    v_cpu = value.detach().to("cpu").contiguous()
-                    dtype_str = self._dtype_to_name(v_cpu.dtype)
-                    data_bytes = v_cpu.numpy().tobytes()
-                    entry = AdditionalInformationEntry(
-                        tensor_data=data_bytes,
-                        tensor_shape=[int(x) for x in list(v_cpu.shape)],
-                        tensor_dtype=dtype_str,
-                    )
-                elif isinstance(value, list):
-                    entry = AdditionalInformationEntry(list_data=value)
-                else:
-                    raise ValueError("additional_information values must be Tensor or list")
-                entries[key] = entry
-            additional_information_payload = AdditionalInformationPayload(entries=entries)
-
-        return OmniEngineCoreRequest(
-            request_id=request_id,
-            prompt_token_ids=prompt_token_ids,
-            mm_features=mm_features,
-            sampling_params=sampling_params,
-            pooling_params=pooling_params,
-            arrival_time=arrival_time,
-            lora_request=lora_request,
-            cache_salt=decoder_inputs.get("cache_salt"),
-            priority=priority,
-            data_parallel_rank=data_parallel_rank,
-            trace_headers=trace_headers,
-            prompt_embeds=prompt_embeds,
-            additional_information=additional_information_payload,
-            resumable=resumable,
-        )
-
-    @staticmethod
-    def _decode_prompt_embeds(payload: PromptEmbedsPayload) -> torch.Tensor:
-        dtype = getattr(np, payload.dtype)
-        arr = np.frombuffer(payload.data, dtype=dtype)
-        arr = arr.reshape(payload.shape)
-        return torch.from_numpy(arr)

vllm_omni/engine/orchestrator.py

@@ -0,0 +1,812 @@
+"""
+Orchestrator for vLLM-Omni multi-stage runtime.
+
+Runs inside a background thread with its own asyncio event loop.
+Owns all StageEngineCoreClient instances, input/output processors,
+and handles stage-to-stage transfer logic.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import copy
+import time as _time
+from dataclasses import dataclass, field
+from typing import Any
+
+import janus
+import torch
+from vllm.config import ModelConfig
+from vllm.logger import init_logger
+from vllm.outputs import RequestOutput
+from vllm.pooling_params import PoolingParams
+from vllm.sampling_params import SamplingParams
+from vllm.v1.engine import EngineCoreOutputs
+
+from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length
+from vllm_omni.engine import (
+    OmniEngineCoreRequest,
+)
+from vllm_omni.engine.serialization import serialize_additional_information
+from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics
+from vllm_omni.metrics.stats import StageStats
+from vllm_omni.metrics.utils import count_tokens_from_outputs
+
+logger = init_logger(__name__)
+
+
+def build_engine_core_request_from_tokens(
+    request_id: str,
+    prompt: dict[str, Any],
+    params: SamplingParams | PoolingParams,
+    arrival_time: float | None = None,
+    model_config: ModelConfig | None = None,
+) -> OmniEngineCoreRequest:
+    """Build an OmniEngineCoreRequest directly from an OmniTokensPrompt.
+
+    Lightweight alternative to the full InputProcessor pipeline - skips
+    tokenization, multimodal preprocessing, LoRA validation, and platform
+    validation.  Intended for stage 1+ where the upstream stage has already
+    produced token IDs and optional embeddings.
+    """
+    if arrival_time is None:
+        arrival_time = _time.time()
+
+    prompt_token_ids = prompt["prompt_token_ids"]
+
+    # Clone params and set max_tokens if needed
+    sampling_params = None
+    pooling_params = None
+    if isinstance(params, SamplingParams):
+        sampling_params = params.clone()
+        if sampling_params.max_tokens is None and model_config is not None:
+            sampling_params.max_tokens = model_config.max_model_len - len(prompt_token_ids)
+    else:
+        pooling_params = params.clone()
+
+    prompt_embeds: torch.Tensor | None = prompt.get("prompt_embeds")
+
+    # Serialize additional_information if present
+    additional_info_payload = serialize_additional_information(
+        prompt.get("additional_information"),
+        log_prefix=f"build_engine_core_request_from_tokens req={request_id}",
+    )
+
+    return OmniEngineCoreRequest(
+        request_id=request_id,
+        prompt_token_ids=prompt_token_ids,
+        mm_features=None,
+        sampling_params=sampling_params,
+        pooling_params=pooling_params,
+        arrival_time=arrival_time,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+        prompt_embeds=prompt_embeds,
+        additional_information=additional_info_payload,
+    )
+
+
+# ============================================================
+# Orchestrator internals (run inside the background thread)
+# ============================================================
+
+
+@dataclass
+class OrchestratorRequestState:
+    """Per-request bookkeeping inside the Orchestrator."""
+
+    request_id: str
+    prompt: Any = None
+    sampling_params_list: list[Any] = field(default_factory=list)
+    final_stage_id: int = -1
+
+    # Metrics: timestamp when request was submitted to each stage
+    stage_submit_ts: dict[int, float] = field(default_factory=dict)
+
+
+class Orchestrator:
+    """Runs inside a background thread's asyncio event loop.
+
+    Owns all StageEngineCoreClient instances, input/output processors,
+    and handles stage-to-stage transfer logic.
+    """
+
+    def __init__(
+        self,
+        request_async_queue: janus.AsyncQueue[dict[str, Any]],
+        output_async_queue: janus.AsyncQueue[dict[str, Any]],
+        rpc_async_queue: janus.AsyncQueue[dict[str, Any]],
+        stage_clients: list[Any],
+        output_processors: list[Any],
+        stage_vllm_configs: list[Any],
+        *,
+        async_chunk: bool = False,
+    ) -> None:
+        self.request_async_queue = request_async_queue
+        self.output_async_queue = output_async_queue
+        self.rpc_async_queue = rpc_async_queue
+
+        self.num_stages = len(stage_clients)
+        self.async_chunk = bool(async_chunk)
+
+        self.stage_clients: list[Any] = stage_clients
+        self.output_processors: list[Any] = output_processors
+        self.stage_vllm_configs: list[Any] = stage_vllm_configs
+
+        # Per-request state
+        self.request_states: dict[str, OrchestratorRequestState] = {}
+
+        # CFG companion tracking
+        self._companion_map: dict[str, dict[str, str]] = {}
+        self._companion_to_parent: dict[str, str] = {}
+        self._companion_ids: set[str] = set()
+        self._companion_done: dict[str, set[str]] = {}
+        self._deferred_parents: dict[str, dict[str, Any]] = {}
+
+        # Per-stage metrics accumulators.
+        self._batch_seq: list[int] = [0] * self.num_stages
+        self._agg_total_tokens: list[int] = [0] * self.num_stages
+        self._agg_total_gen_time_ms: list[float] = [0.0] * self.num_stages
+
+        # Shutdown coordination
+        self._shutdown_event = asyncio.Event()
+        self._stages_shutdown = False
+
+    async def run(self) -> None:
+        """Main entry point for the Orchestrator event loop."""
+        logger.info("[Orchestrator] Starting event loop")
+
+        request_task = asyncio.create_task(self._request_handler(), name="orchestrator-request-handler")
+        output_task = asyncio.create_task(
+            self._orchestration_output_handler(),
+            name="orchestrator-stage-output-handler",
+        )
+
+        try:
+            # Run both tasks concurrently; if either fails the other is cancelled.
+            await asyncio.gather(request_task, output_task)
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            logger.exception("[Orchestrator] Fatal error in orchestrator tasks")
+            raise
+        finally:
+            self._shutdown_event.set()
+            for t in (request_task, output_task):
+                if not t.done():
+                    t.cancel()
+            try:
+                await asyncio.gather(request_task, output_task, return_exceptions=True)
+            except Exception:
+                pass
+
+            self._shutdown_stages()
+
+            # Cancel any remaining tasks spawned by wait_for / gather so
+            # the event loop can close cleanly without "pending task" errors.
+            loop = asyncio.get_running_loop()
+            pending = [t for t in asyncio.all_tasks(loop) if t is not asyncio.current_task() and not t.done()]
+            for t in pending:
+                t.cancel()
+            if pending:
+                await asyncio.gather(*pending, return_exceptions=True)
+
+    async def _request_handler(self) -> None:
+        """Read messages from the main thread via request_async_queue."""
+        while True:
+            msg = await self.request_async_queue.get()
+            msg_type = msg.get("type")
+
+            if msg_type == "add_request":
+                await self._handle_add_request(msg)
+            elif msg_type == "add_companion_request":
+                await self._handle_add_companion(msg)
+            elif msg_type == "abort":
+                await self._handle_abort(msg)
+            elif msg_type == "collective_rpc":
+                await self._handle_collective_rpc(msg)
+            elif msg_type == "shutdown":
+                logger.info("[Orchestrator] Received shutdown signal")
+                self._shutdown_event.set()
+                self._shutdown_stages()
+                break
+            else:
+                logger.warning(f"[Orchestrator] Unknown message type: {msg_type}")
+
+    async def _orchestration_output_handler(self) -> None:
+        """Poll all stages, handle transfers, send final outputs to main."""
+        try:
+            await self._orchestration_loop()
+        except asyncio.CancelledError:
+            logger.debug("[Orchestrator] _orchestration_output_handler cancelled")
+            return
+
+    async def _orchestration_loop(self) -> None:
+        """Inner loop for _orchestration_output_handler (clean cancellation).
+
+        Control flow: poll raw → process through output processor → route.
+        """
+        while not self._shutdown_event.is_set():
+            idle = True
+            for stage_id in range(self.num_stages):
+                if self._shutdown_event.is_set():
+                    return
+
+                # 1) Diffusion stage: poll non-blocking queue
+                # TODO (Peiqi): the output of diffusion stage is OmniRequestOutput,
+                # which is different from EngineCoreOutputs (LLM stages). We may want to unify
+                # the output format in the future to simplify the processing logic in Orchestrator.
+                stage_client = self.stage_clients[stage_id]
+                if stage_client.stage_type == "diffusion":
+                    output = stage_client.get_diffusion_output_async()
+                    if output is not None:
+                        idle = False
+                        req_state = self.request_states.get(output.request_id)
+                        if req_state is not None:
+                            stage_metrics = self._build_stage_metrics(stage_id, output.request_id, [output], req_state)
+                            await self._route_output(stage_id, output, req_state, stage_metrics)
+                    continue
+
+                # 1) Poll raw outputs from the stage
+                try:
+                    raw_outputs = await asyncio.wait_for(self._poll_stage_raw(stage_id), timeout=0.001)
+                except asyncio.TimeoutError:
+                    continue
+                except asyncio.CancelledError:
+                    raise
+                except Exception:
+                    if self._shutdown_event.is_set():
+                        return
+                    logger.exception(
+                        "[Orchestrator] _poll_stage_raw failed for stage-%s",
+                        stage_id,
+                    )
+                    raise
+
+                if raw_outputs is None:
+                    continue
+                idle = False
+
+                # 2) Process raw outputs through the output processor
+                request_outputs = await self._process_stage_outputs(stage_id, raw_outputs)
+
+                # 3) Route each processed output
+                for output in request_outputs:
+                    req_state = self.request_states.get(output.request_id)
+                    if req_state is None:
+                        logger.warning(
+                            "[Orchestrator] Dropping output for unknown req %s at stage-%s (known reqs: %s)",
+                            output.request_id,
+                            stage_id,
+                            list(self.request_states.keys()),
+                        )
+                        continue
+                    stage_metrics = None
+                    if output.finished:
+                        stage_metrics = self._build_stage_metrics(
+                            stage_id,
+                            output.request_id,
+                            [output],
+                            req_state,
+                        )
+                    await self._route_output(stage_id, output, req_state, stage_metrics)
+
+            if idle:
+                await asyncio.sleep(0.001)
+            else:
+                await asyncio.sleep(0)
+
+    async def _route_output(
+        self,
+        stage_id: int,
+        output: Any,
+        req_state: OrchestratorRequestState,
+        stage_metrics: Any,
+    ) -> None:
+        """Route a processed output: send to main thread and/or forward to next stage."""
+        req_id = output.request_id
+        finished = output.finished
+        submit_ts = req_state.stage_submit_ts.get(stage_id)
+        stage_client = self.stage_clients[stage_id]
+
+        # CFG companion handling: companions don't produce user-visible output
+        # and don't forward to the next stage directly.
+        if finished and req_id in self._companion_ids:
+            parent_id = self._companion_to_parent.get(req_id)
+            if parent_id is not None:
+                self._companion_done.setdefault(parent_id, set()).add(req_id)
+                logger.debug(
+                    "[Orchestrator] CFG companion %s done (parent=%s)",
+                    req_id,
+                    parent_id,
+                )
+                # Check if parent is waiting and all companions are done
+                if parent_id in self._deferred_parents and self._all_companions_done(parent_id):
+                    deferred = self._deferred_parents.pop(parent_id)
+                    parent_state = self.request_states.get(parent_id)
+                    if parent_state is not None:
+                        await self._forward_to_next_stage(
+                            parent_id,
+                            deferred["stage_id"],
+                            deferred["output"],
+                            parent_state,
+                        )
+            self.request_states.pop(req_id, None)
+            return
+
+        if stage_client.final_output:
+            await self.output_async_queue.put(
+                {
+                    "type": "output",
+                    "request_id": req_id,
+                    "stage_id": stage_id,
+                    "engine_outputs": output,
+                    "metrics": stage_metrics,
+                    "finished": finished and stage_id == req_state.final_stage_id,
+                    "stage_submit_ts": submit_ts,
+                }
+            )
+        elif stage_metrics is not None:
+            await self.output_async_queue.put(
+                {
+                    "type": "stage_metrics",
+                    "request_id": req_id,
+                    "stage_id": stage_id,
+                    "metrics": stage_metrics,
+                    "stage_submit_ts": submit_ts,
+                }
+            )
+
+        if finished and stage_id < req_state.final_stage_id and not self.async_chunk:
+            # If this parent has CFG companions, defer forwarding until all done
+            if req_id in self._companion_map and not self._all_companions_done(req_id):
+                self._deferred_parents[req_id] = {
+                    "stage_id": stage_id,
+                    "output": output,
+                }
+                logger.debug(
+                    "[Orchestrator] Parent %s deferred, waiting for CFG companions",
+                    req_id,
+                )
+            else:
+                await self._forward_to_next_stage(req_id, stage_id, output, req_state)
+
+        if finished and stage_id == req_state.final_stage_id:
+            self._cleanup_companion_state(req_id)
+            self.request_states.pop(req_id, None)
+
+    def _cleanup_companion_state(self, parent_id: str) -> None:
+        """Remove all companion tracking state for a completed parent."""
+        role_map = self._companion_map.pop(parent_id, {})
+        for cid in role_map.values():
+            self._companion_ids.discard(cid)
+            self._companion_to_parent.pop(cid, None)
+        self._companion_done.pop(parent_id, None)
+        self._deferred_parents.pop(parent_id, None)
+
+    def _all_companions_done(self, parent_id: str) -> bool:
+        """Check whether all CFG companions for a parent request have finished."""
+        role_map = self._companion_map.get(parent_id, {})
+        if not role_map:
+            return True
+        done_set = self._companion_done.get(parent_id, set())
+        return all(cid in done_set for cid in role_map.values())
+
+    def _build_stage_metrics(
+        self,
+        stage_id: int,
+        req_id: str,
+        request_outputs: list[RequestOutput],
+        req_state: OrchestratorRequestState,
+    ) -> StageRequestMetrics:
+        """Build StageRequestMetrics for a finished request at a stage.
+
+        Reuses StageRequestMetrics so OrchestratorMetrics and downstream
+        metric handlers can consume a stable schema.
+        """
+        now = _time.time()
+        submit_ts = req_state.stage_submit_ts.get(stage_id, now)
+        stage_gen_time_ms = (now - submit_ts) * 1000.0
+
+        num_tokens_out = count_tokens_from_outputs(request_outputs)
+        num_tokens_in = 0
+        if stage_id == 0:
+            for ro in request_outputs:
+                ptids = getattr(ro, "prompt_token_ids", None)
+                if ptids is not None:
+                    num_tokens_in += len(ptids)
+
+        # Monotonic batch counter per stage.
+        self._batch_seq[stage_id] += 1
+        batch_id = self._batch_seq[stage_id]
+
+        # Accumulate for running-average stage_stats
+        self._agg_total_tokens[stage_id] += num_tokens_out
+        self._agg_total_gen_time_ms[stage_id] += stage_gen_time_ms
+
+        return StageRequestMetrics(
+            num_tokens_in=num_tokens_in,
+            num_tokens_out=num_tokens_out,
+            stage_gen_time_ms=stage_gen_time_ms,
+            batch_id=batch_id,
+            batch_size=1,
+            rx_decode_time_ms=0.0,
+            rx_transfer_bytes=0,
+            rx_in_flight_time_ms=0.0,
+            stage_stats=StageStats(
+                total_token=self._agg_total_tokens[stage_id],
+                total_gen_time_ms=self._agg_total_gen_time_ms[stage_id],
+            ),
+        )
+
+    async def _forward_to_next_stage(
+        self,
+        req_id: str,
+        stage_id: int,
+        output: Any,
+        req_state: OrchestratorRequestState,
+    ) -> None:
+        """Forward output from current stage to the next stage.
+
+        Handles the full pipeline: set outputs on current stage, compute
+        next-stage inputs, build lightweight requests, and submit them.
+        """
+        next_stage_id = stage_id + 1
+        next_client = self.stage_clients[next_stage_id]
+        params = req_state.sampling_params_list[next_stage_id]
+
+        if next_client.stage_type == "diffusion":
+            self.stage_clients[stage_id].set_engine_outputs([output])
+            if next_client.custom_process_input_func is not None:
+                diffusion_prompt = next_client.custom_process_input_func(
+                    self.stage_clients,
+                    next_client.engine_input_source,
+                    req_state.prompt,
+                    False,
+                )
+                if isinstance(diffusion_prompt, list):
+                    diffusion_prompt = diffusion_prompt[0]
+            else:
+                diffusion_prompt = req_state.prompt
+
+            # Attach CFG companion KV request IDs so the diffusion model
+            # runner can fetch companion KV caches alongside the primary one.
+            cfg_ids = self._companion_map.get(req_id)
+            if cfg_ids:
+                from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+                if isinstance(params, OmniDiffusionSamplingParams):
+                    params = copy.deepcopy(params)
+                    params.cfg_kv_request_ids = cfg_ids
+                    logger.info(
+                        "[Orchestrator] Attaching cfg_kv_request_ids=%s to req %s",
+                        cfg_ids,
+                        req_id,
+                    )
+
+            await next_client.add_request_async(req_id, diffusion_prompt, params)
+            req_state.stage_submit_ts[next_stage_id] = _time.time()
+            return
+
+        self.stage_clients[stage_id].set_engine_outputs([output])
+
+        # Process inputs for next stage
+        try:
+            next_inputs = next_client.process_engine_inputs(
+                stage_list=self.stage_clients,
+                prompt=req_state.prompt,
+            )
+        except Exception:
+            logger.exception(
+                "[Orchestrator] req=%s process_engine_inputs FAILED for stage-%s",
+                req_id,
+                next_stage_id,
+            )
+            raise
+
+        # Build and submit requests for each input
+        for next_input in next_inputs:
+            request = build_engine_core_request_from_tokens(
+                request_id=req_id,
+                prompt=next_input,
+                params=params,
+                model_config=self.stage_vllm_configs[next_stage_id].model_config,
+            )
+
+            # TODO: Here we directly use the req id to assign.
+            request.external_req_id = request.request_id
+
+            self.output_processors[next_stage_id].add_request(
+                request=request,
+                prompt=None,
+                parent_req=None,
+                request_index=0,
+                queue=None,
+            )
+
+            await next_client.add_request_async(request)
+
+        # Record submit timestamp for the next stage
+        req_state.stage_submit_ts[next_stage_id] = _time.time()
+
+    async def _poll_stage_raw(self, stage_id: int) -> EngineCoreOutputs | None:
+        """Pull raw EngineCoreOutputs from a stage client without processing.
+
+        Returns the raw outputs object, or None when there is nothing
+        to consume.
+        """
+        outputs = await self.stage_clients[stage_id].get_output_async()
+        if not outputs.outputs:
+            return None
+        return outputs
+
+    async def _process_stage_outputs(self, stage_id: int, raw_outputs: EngineCoreOutputs) -> list[RequestOutput]:
+        """Run the output processor on raw outputs, returning RequestOutputs.
+
+        Also handles abort forwarding and scheduler stats updates.
+        """
+        processor = self.output_processors[stage_id]
+
+        processed = processor.process_outputs(
+            raw_outputs.outputs,
+            raw_outputs.timestamp,
+            None,
+        )
+
+        if processed.reqs_to_abort:
+            await self.stage_clients[stage_id].abort_requests_async(processed.reqs_to_abort)
+
+        if raw_outputs.scheduler_stats is not None:
+            processor.update_scheduler_stats(raw_outputs.scheduler_stats)
+
+        return processed.request_outputs
+
+    async def _handle_add_request(self, msg: dict[str, Any]) -> None:
+        """Handle an add_request message from the main thread."""
+        stage_id = 0
+        request_id = msg["request_id"]
+        prompt = msg["prompt"]
+        original_prompt = msg.get("original_prompt", prompt)
+        sampling_params_list = msg["sampling_params_list"]
+        if not sampling_params_list:
+            raise ValueError(f"Missing sampling params for stage 0. Got {len(sampling_params_list)} stage params.")
+        params = sampling_params_list[0]
+        final_stage_id = msg["final_stage_id"]
+
+        logger.info(
+            "[Orchestrator] _handle_add_request: stage=%s req=%s "
+            "prompt_type=%s original_prompt_type=%s final_stage=%s "
+            "num_sampling_params=%d",
+            stage_id,
+            request_id,
+            type(prompt).__name__,
+            type(original_prompt).__name__,
+            final_stage_id,
+            len(sampling_params_list),
+        )
+
+        # Track request state - use original_prompt so downstream stages
+        # (e.g. thinker2talker) can access the raw dict with multi_modal_data.
+        req_state = OrchestratorRequestState(
+            request_id=request_id,
+            prompt=original_prompt,
+            sampling_params_list=sampling_params_list,
+            final_stage_id=final_stage_id,
+        )
+        req_state.stage_submit_ts[stage_id] = _time.time()
+        self.request_states[request_id] = req_state
+
+        # Stage-0 prompt is already a fully-formed OmniEngineCoreRequest
+        # (pre-processed by AsyncOmniEngine.add_request, output processor
+        # already registered there) - submit directly.
+        request = prompt
+        stage_client = self.stage_clients[stage_id]
+        if stage_client.stage_type == "diffusion":
+            await stage_client.add_request_async(request_id, prompt, params)
+        else:
+            await stage_client.add_request_async(request)
+
+        if self.async_chunk and stage_id == 0 and final_stage_id > 0:
+            await self._prewarm_async_chunk_stages(request_id, request, req_state)
+
+    async def _prewarm_async_chunk_stages(
+        self,
+        request_id: str,
+        stage0_request: Any,
+        req_state: OrchestratorRequestState,
+    ) -> None:
+        """Pre-submit downstream stages for async-chunk mode.
+
+        In async-chunk mode, stages exchange data through connectors/chunk adapters,
+        so downstream stages should be armed once at request start instead of waiting
+        for stage-finished forwarding.
+        """
+        if req_state.final_stage_id <= 0:
+            return
+
+        prompt_token_ids = getattr(stage0_request, "prompt_token_ids", None)
+        if prompt_token_ids is None:
+            logger.warning(
+                "[Orchestrator] async_chunk prewarm skipped for req=%s: stage0 prompt_token_ids missing",
+                request_id,
+            )
+            return
+
+        # Pre-arm stage-1+ with placeholder prompt IDs.
+        try:
+            next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
+        except Exception:
+            next_prompt_len = max(1, len(prompt_token_ids))
+        original_prompt = req_state.prompt
+        if isinstance(original_prompt, dict):
+            base_input = copy.deepcopy(original_prompt)
+        else:
+            base_input = {}
+        base_input["prompt_token_ids"] = [0] * next_prompt_len
+        base_input["multi_modal_data"] = None
+        base_input["mm_processor_kwargs"] = None
+
+        for next_stage_id in range(1, req_state.final_stage_id + 1):
+            next_client = self.stage_clients[next_stage_id]
+            params = req_state.sampling_params_list[next_stage_id]
+
+            if next_client.stage_type == "diffusion":
+                await next_client.add_request_async(request_id, req_state.prompt, params)
+                req_state.stage_submit_ts[next_stage_id] = _time.time()
+                continue
+
+            request = build_engine_core_request_from_tokens(
+                request_id=request_id,
+                prompt=base_input,
+                params=params,
+                model_config=self.stage_vllm_configs[next_stage_id].model_config,
+            )
+            request.external_req_id = request.request_id
+
+            self.output_processors[next_stage_id].add_request(
+                request=request,
+                prompt=None,
+                parent_req=None,
+                request_index=0,
+                queue=None,
+            )
+            await next_client.add_request_async(request)
+            req_state.stage_submit_ts[next_stage_id] = _time.time()
+
+    async def _handle_add_companion(self, msg: dict[str, Any]) -> None:
+        """Handle an add_companion_request message: submit companion to stage 0."""
+        companion_id = msg["companion_id"]
+        parent_id = msg["parent_id"]
+        role = msg["role"]
+        companion_prompt = msg["prompt"]
+        sampling_params_list = msg["sampling_params_list"]
+
+        # Register companion mapping
+        if parent_id not in self._companion_map:
+            self._companion_map[parent_id] = {}
+        self._companion_map[parent_id][role] = companion_id
+        self._companion_ids.add(companion_id)
+        self._companion_to_parent[companion_id] = parent_id
+        self._companion_done.setdefault(parent_id, set())
+
+        companion_state = OrchestratorRequestState(
+            request_id=companion_id,
+            prompt=companion_prompt,
+            sampling_params_list=sampling_params_list,
+            final_stage_id=0,
+        )
+        companion_state.stage_submit_ts[0] = _time.time()
+        self.request_states[companion_id] = companion_state
+
+        request = companion_prompt  # Already a processed OmniEngineCoreRequest
+        stage_client = self.stage_clients[0]
+        await stage_client.add_request_async(request)
+
+        logger.info(
+            "[Orchestrator] CFG companion submitted: %s (role=%s, parent=%s)",
+            companion_id,
+            role,
+            parent_id,
+        )
+
+    async def _handle_abort(self, msg: dict[str, Any]) -> None:
+        """Handle an abort message from the main thread."""
+        request_ids = msg["request_ids"]
+        # Also abort any CFG companions for aborted parents
+        companion_ids_to_abort: list[str] = []
+        for req_id in request_ids:
+            role_map = self._companion_map.pop(req_id, {})
+            for cid in role_map.values():
+                companion_ids_to_abort.append(cid)
+                self._companion_ids.discard(cid)
+                self._companion_to_parent.pop(cid, None)
+                self.request_states.pop(cid, None)
+            self._companion_done.pop(req_id, None)
+            self._deferred_parents.pop(req_id, None)
+
+        all_ids_to_abort = list(request_ids) + companion_ids_to_abort
+        for stage_id in range(self.num_stages):
+            await self.stage_clients[stage_id].abort_requests_async(all_ids_to_abort)
+        for req_id in request_ids:
+            self.request_states.pop(req_id, None)
+        logger.info("[Orchestrator] Aborted request(s) %s", request_ids)
+
+    async def _handle_collective_rpc(self, msg: dict[str, Any]) -> None:
+        """Handle a control-plane RPC request from the main thread.
+
+        TODO(AsyncOmni): parallelize stage dispatch if control latency becomes
+        noticeable. The current sequential fanout keeps the first version simple
+        and deterministic.
+        """
+        rpc_id = msg["rpc_id"]
+        method = msg["method"]
+        timeout = msg.get("timeout")
+        args = tuple(msg.get("args", ()))
+        kwargs = dict(msg.get("kwargs") or {})
+        requested_stage_ids = msg.get("stage_ids")
+        stage_ids = list(range(self.num_stages)) if requested_stage_ids is None else list(requested_stage_ids)
+
+        results: list[Any] = []
+        for stage_id in stage_ids:
+            if stage_id < 0 or stage_id >= self.num_stages:
+                results.append(
+                    {
+                        "supported": False,
+                        "todo": True,
+                        "error": f"Invalid stage id {stage_id}",
+                    }
+                )
+                continue
+
+            stage_client = self.stage_clients[stage_id]
+            try:
+                if hasattr(stage_client, "collective_rpc_async"):
+                    stage_result = await stage_client.collective_rpc_async(
+                        method=method,
+                        timeout=timeout,
+                        args=args,
+                        kwargs=kwargs,
+                    )
+                else:
+                    stage_result = {
+                        "supported": False,
+                        "todo": True,
+                        "reason": (f"{stage_client.__class__.__name__}.collective_rpc_async is not implemented yet"),
+                    }
+            except Exception as exc:
+                logger.exception(
+                    "[Orchestrator] collective_rpc failed: stage=%s method=%s",
+                    stage_id,
+                    method,
+                )
+                stage_result = {
+                    "supported": False,
+                    "error": str(exc),
+                }
+
+            results.append(stage_result)
+
+        await self.rpc_async_queue.put(
+            {
+                "type": "collective_rpc_result",
+                "rpc_id": rpc_id,
+                "method": method,
+                "stage_ids": stage_ids,
+                "results": results,
+            }
+        )
+
+    def _shutdown_stages(self) -> None:
+        """Shutdown all stage clients."""
+        if self._stages_shutdown:
+            return
+
+        self._stages_shutdown = True
+        logger.info("[Orchestrator] Shutting down all stages")
+        for stage_id, stage_client in enumerate(self.stage_clients):
+            try:
+                stage_client.shutdown()
+                logger.info(f"[Orchestrator] Stage {stage_id} shut down")
+            except Exception as e:
+                logger.warning(f"[Orchestrator] Failed to shutdown stage {stage_id}: {e}")

vllm_omni/engine/serialization.py

@@ -0,0 +1,81 @@
+"""Shared serialization helpers for omni engine request payloads."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.engine import (
+    AdditionalInformationEntry,
+    AdditionalInformationPayload,
+)
+
+logger = init_logger(__name__)
+
+
+def dtype_to_name(dtype: torch.dtype) -> str:
+    """Convert torch dtype to a stable string name for serialization."""
+    mapping = {
+        torch.float32: "float32",
+        torch.float: "float32",
+        torch.float16: "float16",
+        torch.half: "float16",
+        torch.bfloat16: "bfloat16",
+        torch.float64: "float64",
+        torch.double: "float64",
+        torch.int64: "int64",
+        torch.long: "int64",
+        torch.int32: "int32",
+        torch.int: "int32",
+        torch.int16: "int16",
+        torch.short: "int16",
+        torch.int8: "int8",
+        torch.uint8: "uint8",
+        torch.bool: "bool",
+    }
+    return mapping.get(dtype, str(dtype).replace("torch.", ""))
+
+
+def serialize_additional_information(
+    raw_info: dict[str, Any] | AdditionalInformationPayload | None,
+    *,
+    log_prefix: str | None = None,
+) -> AdditionalInformationPayload | None:
+    """Serialize omni request metadata for EngineCore transport."""
+    if raw_info is None:
+        return None
+    if isinstance(raw_info, AdditionalInformationPayload):
+        return raw_info
+
+    entries: dict[str, AdditionalInformationEntry] = {}
+    for key, value in raw_info.items():
+        if isinstance(value, torch.Tensor):
+            value_cpu = value.detach().to("cpu").contiguous()
+            entries[key] = AdditionalInformationEntry(
+                tensor_data=value_cpu.numpy().tobytes(),
+                tensor_shape=list(value_cpu.shape),
+                tensor_dtype=dtype_to_name(value_cpu.dtype),
+            )
+            continue
+
+        if isinstance(value, list):
+            entries[key] = AdditionalInformationEntry(list_data=value)
+            continue
+
+        if log_prefix is None:
+            logger.warning(
+                "Dropping unsupported additional_information key=%s type=%s",
+                key,
+                type(value).__name__,
+            )
+        else:
+            logger.warning(
+                "[%s] Dropping unsupported additional_information key=%s type=%s",
+                log_prefix,
+                key,
+                type(value).__name__,
+            )
+
+    return AdditionalInformationPayload(entries=entries) if entries else None

vllm_omni/engine/stage_engine_core_client.py

@@ -0,0 +1,168 @@
+"""
+Stage Engine Core Client for vLLM-Omni multi-stage runtime.
+
+Directly inherits from vLLM's AsyncMPClient to reuse EngineCore architecture.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from vllm.logger import init_logger
+from vllm.v1.engine import EngineCoreRequest
+from vllm.v1.engine.core_client import AsyncMPClient
+
+from vllm_omni.engine.stage_init_utils import StageMetadata
+
+if TYPE_CHECKING:
+    from vllm.v1.engine import EngineCoreOutput
+
+    from vllm_omni.inputs.data import OmniTokensPrompt
+
+logger = init_logger(__name__)
+
+
+class StageEngineCoreClient(AsyncMPClient):
+    """Stage async client that inherits from vLLM's AsyncMPClient.
+
+    Fully reuses AsyncMPClient.__init__ for:
+    - ZMQ setup, sockets
+    - launch_core_engines() -> EngineCoreProc
+    - outputs_queue, output_queue_task
+    - All utility methods (shutdown, get_output_async, abort_requests_async, etc.)
+
+    This is the async version of StageMPClient, designed for use with AsyncOmniEngine.
+    """
+
+    def __init__(
+        self,
+        vllm_config: Any,
+        executor_class: type,
+        metadata: StageMetadata,
+        client_addresses: dict[str, str] | None = None,
+        engine_manager: Any = None,
+        coordinator: Any = None,
+    ):
+        """Create an async EngineCore client for a single stage.
+
+        All heavy init (config extraction, plugin loading, device setup,
+        engine args building, device locking) is done by the Orchestrator
+        via helpers in stage_init_utils.py. This constructor just stores metadata
+        and calls super().__init__().
+        """
+        # -------- Stage metadata (public fields used at runtime) --------
+        self.stage_id = metadata.stage_id
+        self.stage_type = metadata.stage_type
+        self.engine_output_type = metadata.engine_output_type
+        self.is_comprehension = metadata.is_comprehension
+        self.requires_multimodal_data = metadata.requires_multimodal_data
+        self.engine_input_source = metadata.engine_input_source
+        self.final_output = metadata.final_output
+        self.final_output_type = metadata.final_output_type
+        self.default_sampling_params = metadata.default_sampling_params
+        self.custom_process_input_func = metadata.custom_process_input_func
+        self.model_stage = metadata.model_stage
+
+        self.engine_outputs: Any = None
+
+        logger.info(
+            "[StageEngineCoreClient] Stage-%s initializing EngineCore",
+            self.stage_id,
+        )
+        try:
+            super().__init__(
+                vllm_config,
+                executor_class,
+                log_stats=False,
+                client_addresses=client_addresses,
+            )
+            if engine_manager is not None:
+                self.resources.engine_manager = engine_manager
+            if coordinator is not None:
+                self.resources.coordinator = coordinator
+        except Exception:
+            logger.exception(
+                "[StageEngineCoreClient] Stage-%s EngineCore init failed",
+                self.stage_id,
+            )
+            try:
+                self.shutdown()
+            except Exception as shutdown_error:
+                logger.warning(
+                    "[StageEngineCoreClient] Stage-%s cleanup after init failure failed: %s",
+                    self.stage_id,
+                    shutdown_error,
+                )
+            raise
+        logger.info(
+            "[StageEngineCoreClient] Stage-%s EngineCore running",
+            self.stage_id,
+        )
+
+    # ==================== Overrides ====================
+
+    async def add_request_async(self, request: EngineCoreRequest) -> None:
+        """Add request to the stage engine core."""
+        logger.info(f"[StageEngineCoreClient] Stage-{self.stage_id} adding request: {request.request_id}")
+        await super().add_request_async(request)
+
+    # ==================== Stage Methods ====================
+
+    def set_engine_outputs(self, engine_outputs: EngineCoreOutput) -> None:
+        """Set engine outputs (called by orchestrator)."""
+        self.engine_outputs = engine_outputs
+
+    def process_engine_inputs(
+        self,
+        stage_list: list[Any],
+        prompt: OmniTokensPrompt | list[OmniTokensPrompt] | None = None,
+    ) -> list[OmniTokensPrompt]:
+        """Process inputs from upstream stages."""
+        from vllm_omni.inputs.data import OmniTokensPrompt
+
+        if self.custom_process_input_func is not None:
+            return self.custom_process_input_func(
+                stage_list,
+                self.engine_input_source,
+                prompt,
+                self.requires_multimodal_data,
+            )
+
+        if not self.engine_input_source:
+            raise ValueError(f"engine_input_source empty for stage {self.stage_id}")
+
+        source_id = self.engine_input_source[0]
+        source_outputs = stage_list[source_id].engine_outputs
+
+        if not isinstance(prompt, list):
+            prompt = [prompt]
+
+        mm_data = {so.request_id: p.get("multi_modal_data") for so, p in zip(source_outputs, prompt)}
+
+        return [
+            OmniTokensPrompt(
+                prompt_token_ids=so.outputs[0].token_ids,
+                multi_modal_data=(mm_data[so.request_id] if self.requires_multimodal_data else None),
+            )
+            for so in source_outputs
+        ]
+
+    async def collective_rpc_async(
+        self,
+        method: str,
+        timeout: float | None = None,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> Any:
+        """Forward control RPCs to the underlying AsyncMPClient stage engine.
+
+        Each ``StageEngineCoreClient`` already represents one logical stage, so
+        stage-scoped control operations should be executed here and then fanned
+        in-core across the workers managed by this EngineCore client.
+        """
+        return await super().collective_rpc_async(
+            method=method,
+            timeout=timeout,
+            args=args,
+            kwargs=kwargs,
+        )

vllm_omni/engine/stage_init_utils.py

@@ -0,0 +1,525 @@
+"""
+Stage initialization helpers for vLLM-Omni multi-stage runtime.
+
+Extracts orchestration-level init logic (config extraction, plugin loading,
+multiprocessing setup, device mapping, device locking, engine args building)
+out of StageEngineCoreClient into reusable functions.
+"""
+
+from __future__ import annotations
+
+import fcntl
+import importlib
+import multiprocessing as mp
+import os
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Literal
+
+from vllm.logger import init_logger
+from vllm.sampling_params import SamplingParams
+from vllm.usage.usage_lib import UsageContext
+from vllm.v1.engine.input_processor import InputProcessor
+from vllm.v1.executor import Executor
+
+from vllm_omni.engine.arg_utils import OmniEngineArgs
+from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices
+from vllm_omni.entrypoints.utils import filter_dataclass_kwargs, resolve_model_config_path
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams
+
+logger = init_logger(__name__)
+
+
+def _resolve_model_to_local_path(model: str) -> str:
+    """Resolve an HF Hub model ID to a local cache path."""
+    if os.path.isdir(model):
+        return model
+
+    try:
+        from huggingface_hub import snapshot_download
+
+        # Keep init path resolution offline-friendly.
+        return snapshot_download(model, local_files_only=True)
+    except Exception:
+        logger.warning(
+            "[stage_init] Could not resolve %s to local snapshot; using as-is",
+            model,
+        )
+        return model
+
+
+def _resolve_model_tokenizer_paths(model: str, engine_args: dict[str, Any]) -> str:
+    """Apply model_subdir/tokenizer_subdir indirections from stage engine args."""
+    model_subdir = engine_args.pop("model_subdir", None)
+    tokenizer_subdir = engine_args.pop("tokenizer_subdir", None)
+    if model_subdir is None and tokenizer_subdir is None:
+        return model
+
+    resolved_base = _resolve_model_to_local_path(model)
+
+    if model_subdir:
+        model = os.path.join(resolved_base, model_subdir)
+        logger.info("[stage_init] Using model subdirectory: %s", model)
+
+    if tokenizer_subdir is not None:
+        tokenizer_path = os.path.join(resolved_base, tokenizer_subdir) if tokenizer_subdir else resolved_base
+        engine_args["tokenizer"] = tokenizer_path
+        logger.info("[stage_init] Using tokenizer from: %s", tokenizer_path)
+    elif model_subdir and "tokenizer" not in engine_args:
+        # Keep legacy behavior: model in subdir, tokenizer defaults to base path.
+        engine_args["tokenizer"] = resolved_base
+        logger.info("[stage_init] Using tokenizer from base model path: %s", resolved_base)
+
+    return model
+
+
+def resolve_worker_cls(engine_args: dict[str, Any]) -> None:
+    """Resolve worker_cls from worker_type for non-diffusion stages."""
+    worker_type = engine_args.get("worker_type", None)
+    if not worker_type:
+        return
+    worker_cls = engine_args.get("worker_cls")
+    if worker_cls is not None and worker_cls != "auto":
+        return
+
+    from vllm_omni.platforms import current_omni_platform
+
+    worker_type = str(worker_type).lower()
+    if worker_type == "ar":
+        engine_args["worker_cls"] = current_omni_platform.get_omni_ar_worker_cls()
+    elif worker_type == "generation":
+        engine_args["worker_cls"] = current_omni_platform.get_omni_generation_worker_cls()
+    else:
+        raise ValueError(f"Unknown worker_type: {worker_type}")
+
+
+@dataclass
+class StageMetadata:
+    """Lightweight stage attributes extracted from stage_config."""
+
+    stage_id: int
+    stage_type: Literal["llm", "diffusion"]
+    engine_output_type: str | None
+    is_comprehension: bool
+    requires_multimodal_data: bool
+    engine_input_source: list[int]
+    final_output: bool
+    final_output_type: str | None
+    default_sampling_params: OmniSamplingParams
+    custom_process_input_func: Callable | None
+    model_stage: str | None
+    runtime_cfg: Any
+    prompt_expand_func: Callable | None = None
+    cfg_kv_collect_func: Callable | None = None
+
+
+@dataclass
+class StartedLlmStage:
+    """Resources for an LLM stage that has completed startup."""
+
+    stage_id: int
+    metadata: Any
+    vllm_config: Any
+    executor_class: type
+    engine_manager: Any
+    coordinator: Any
+    addresses: Any
+
+
+def extract_stage_metadata(stage_config: Any) -> StageMetadata:
+    """Pure data extraction from a stage_config object."""
+    stage_id: int = stage_config.stage_id
+    stage_type: Literal["llm", "diffusion"] = getattr(stage_config, "stage_type", "llm")
+    engine_args = stage_config.engine_args
+    runtime_cfg = getattr(stage_config, "runtime", {})
+    engine_input_source: list[int] = getattr(stage_config, "engine_input_source", [])
+    final_output: bool = getattr(stage_config, "final_output", False)
+    final_output_type: str | None = getattr(stage_config, "final_output_type", None)
+
+    default_sp = _to_dict(getattr(stage_config, "default_sampling_params", {}))
+    SPClass = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams
+    default_sampling_params: OmniSamplingParams = SPClass(**default_sp)
+
+    custom_process_input_func: Callable | None = None
+    if hasattr(stage_config, "custom_process_input_func"):
+        mod_path, fn_name = stage_config.custom_process_input_func.rsplit(".", 1)
+        custom_process_input_func = getattr(importlib.import_module(mod_path), fn_name)
+
+    prompt_expand_func: Callable | None = None
+    _pef_path = getattr(stage_config, "prompt_expand_func", None)
+    if _pef_path:
+        _mod, _fn = _pef_path.rsplit(".", 1)
+        prompt_expand_func = getattr(importlib.import_module(_mod), _fn)
+
+    cfg_kv_collect_func: Callable | None = None
+    _ckf_path = getattr(stage_config, "cfg_kv_collect_func", None)
+    if _ckf_path:
+        _mod, _fn = _ckf_path.rsplit(".", 1)
+        cfg_kv_collect_func = getattr(importlib.import_module(_mod), _fn)
+
+    if stage_type == "diffusion":
+        return StageMetadata(
+            stage_id=stage_id,
+            stage_type="diffusion",
+            engine_output_type=None,
+            is_comprehension=False,
+            requires_multimodal_data=False,
+            engine_input_source=engine_input_source,
+            final_output=final_output,
+            final_output_type=final_output_type,
+            default_sampling_params=default_sampling_params,
+            custom_process_input_func=custom_process_input_func,
+            model_stage=None,
+            runtime_cfg=runtime_cfg,
+            cfg_kv_collect_func=cfg_kv_collect_func,
+        )
+
+    model_stage = getattr(engine_args, "model_stage", None)
+    engine_output_type = getattr(engine_args, "engine_output_type", None)
+    is_comprehension = getattr(stage_config, "is_comprehension", False)
+    requires_multimodal_data = getattr(runtime_cfg, "requires_multimodal_data", False)
+
+    return StageMetadata(
+        stage_id=stage_id,
+        stage_type=stage_type,
+        engine_output_type=engine_output_type,
+        is_comprehension=is_comprehension,
+        requires_multimodal_data=requires_multimodal_data,
+        engine_input_source=engine_input_source,
+        final_output=final_output,
+        final_output_type=final_output_type,
+        default_sampling_params=default_sampling_params,
+        custom_process_input_func=custom_process_input_func,
+        model_stage=model_stage,
+        runtime_cfg=runtime_cfg,
+        prompt_expand_func=prompt_expand_func,
+    )
+
+
+def prepare_engine_environment() -> None:
+    """One-time global setup: load plugins, set multiprocessing spawn method."""
+    from vllm_omni.plugins import load_omni_general_plugins
+
+    load_omni_general_plugins()
+
+    if os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") != "spawn":
+        os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+        logger.info("[stage_init] Set VLLM_WORKER_MULTIPROC_METHOD=spawn")
+    try:
+        mp.set_start_method("spawn", force=True)
+    except RuntimeError:
+        pass
+
+
+def setup_stage_devices(stage_id: int, runtime_cfg: Any) -> None:
+    """Device mapping via set_stage_devices for a single stage."""
+    try:
+        from vllm_omni.platforms import current_omni_platform
+
+        device_type = current_omni_platform.device_type
+        set_stage_devices(
+            stage_id,
+            runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else None,
+            device_type=device_type,
+        )
+        logger.info(
+            "[stage_init] Stage-%s set devices for %s, runtime devices: %s",
+            stage_id,
+            device_type,
+            runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else None,
+        )
+    except Exception as e:
+        logger.warning("Device setup failed for stage %s: %s", stage_id, e)
+
+
+def build_engine_args_dict(
+    stage_config: Any,
+    model: str,
+    stage_connector_spec: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Build the normalized engine args dict for one stage."""
+    engine_args = stage_config.engine_args
+    stage_type = getattr(stage_config, "stage_type", "llm")
+    stage_id = stage_config.stage_id
+
+    engine_args_dict = _to_dict(engine_args)
+    model = _resolve_model_tokenizer_paths(model, engine_args_dict)
+    engine_args_dict["model"] = model
+    # Stage id must come from stage config instead of inherited CLI kwargs
+    # (e.g. `--stage-id` defaulting to None).
+    engine_args_dict["stage_id"] = stage_id
+    if engine_args_dict.get("async_chunk", False):
+        engine_args_dict["stage_connector_spec"] = dict(stage_connector_spec or {})
+
+    if stage_type != "diffusion":
+        resolve_worker_cls(engine_args_dict)
+
+    return engine_args_dict
+
+
+def build_vllm_config(
+    stage_config: Any,
+    model: str,
+    stage_connector_spec: dict[str, Any] | None = None,
+    engine_args_dict: dict[str, Any] | None = None,
+) -> tuple[Any, type]:
+    """Build engine args, then create VllmConfig and executor_class.
+
+    Returns:
+        (vllm_config, executor_class)
+    """
+    if engine_args_dict is None:
+        engine_args_dict = build_engine_args_dict(
+            stage_config,
+            model,
+            stage_connector_spec=stage_connector_spec,
+        )
+
+    filtered_engine_args_dict = filter_dataclass_kwargs(OmniEngineArgs, engine_args_dict)
+    omni_engine_args = OmniEngineArgs(**filtered_engine_args_dict)
+    vllm_config = omni_engine_args.create_engine_config(usage_context=UsageContext.LLM_CLASS)
+    executor_class = Executor.get_class(vllm_config)
+
+    return vllm_config, executor_class
+
+
+def acquire_device_locks(
+    stage_id: int,
+    engine_args_dict: dict[str, Any],
+    stage_init_timeout: int = 300,
+) -> list[int]:
+    """Acquire exclusive file locks on devices needed by this stage.
+
+    Returns list of lock file descriptors that must be released after init.
+    """
+    lock_fds: list[int] = []
+    try:
+        from vllm_omni.platforms import current_omni_platform
+
+        # Get parallel sizes
+        if "parallel_config" in engine_args_dict:
+            pc = engine_args_dict["parallel_config"]
+            tensor_parallel_size = pc.get("tensor_parallel_size", 1)
+            pipeline_parallel_size = pc.get("pipeline_parallel_size", 1)
+            data_parallel_size = pc.get("data_parallel_size", 1)
+            prefill_context_parallel_size = pc.get("prefill_context_parallel_size", 1)
+            sequence_parallel_size = pc.get("sequence_parallel_size", 1)
+            cfg_parallel_size = pc.get("cfg_parallel_size", 1)
+        else:
+            tensor_parallel_size = engine_args_dict.get("tensor_parallel_size", 1)
+            pipeline_parallel_size = engine_args_dict.get("pipeline_parallel_size", 1)
+            data_parallel_size = engine_args_dict.get("data_parallel_size", 1)
+            prefill_context_parallel_size = engine_args_dict.get("prefill_context_parallel_size", 1)
+            sequence_parallel_size = 1
+            cfg_parallel_size = 1
+
+        num_devices_per_stage = (
+            tensor_parallel_size
+            * pipeline_parallel_size
+            * data_parallel_size
+            * prefill_context_parallel_size
+            * sequence_parallel_size
+            * cfg_parallel_size
+        )
+
+        # Get physical device IDs
+        device_control_env = current_omni_platform.device_control_env_var
+        visible_devices_str = os.environ.get(device_control_env)
+        physical_devices: list[int] = []
+
+        if visible_devices_str:
+            try:
+                physical_devices = [int(x.strip()) for x in visible_devices_str.split(",") if x.strip()]
+            except (ValueError, IndexError):
+                pass
+
+        if not physical_devices:
+            num_devices = current_omni_platform.get_device_count()
+            physical_devices = list(range(num_devices))
+
+        num_devices_to_lock = min(num_devices_per_stage, len(physical_devices))
+        devices_to_lock = sorted(physical_devices[:num_devices_to_lock])
+
+        logger.debug(
+            "Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d, SP=%d, CFG=%d; will lock %d devices: %s",
+            tensor_parallel_size,
+            pipeline_parallel_size,
+            data_parallel_size,
+            prefill_context_parallel_size,
+            sequence_parallel_size,
+            cfg_parallel_size,
+            num_devices_to_lock,
+            devices_to_lock,
+        )
+
+        # Acquire locks
+        wait_start = time.time()
+        for device_id in devices_to_lock:
+            lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
+            lock_acquired = False
+
+            while not lock_acquired:
+                try:
+                    lock_fd = os.open(lock_file, os.O_CREAT | os.O_RDWR, 0o644)
+                    try:
+                        fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                        os.ftruncate(lock_fd, 0)
+                        os.write(lock_fd, f"{os.getpid()}\n".encode())
+                        os.fsync(lock_fd)
+                        lock_acquired = True
+                        lock_fds.append(lock_fd)
+                        logger.debug("Acquired exclusive lock for device %s", device_id)
+                    except BlockingIOError:
+                        os.close(lock_fd)
+                        if time.time() - wait_start > stage_init_timeout:
+                            logger.warning(
+                                "Timeout waiting for device %s initialization lock, proceeding anyway",
+                                device_id,
+                            )
+                            break
+                        time.sleep(0.01)
+                except OSError as e:
+                    logger.debug(
+                        "Failed to acquire lock for device %s: %s, continuing anyway",
+                        device_id,
+                        e,
+                    )
+                    try:
+                        os.close(lock_fd)
+                    except (OSError, NameError):
+                        pass
+                    break
+
+    except Exception as e:
+        logger.debug(
+            "[Stage-%s] Failed to set up sequential initialization lock: %s",
+            stage_id,
+            e,
+        )
+
+    return lock_fds
+
+
+def release_device_locks(lock_fds: list[int]) -> None:
+    """Release file locks acquired by acquire_device_locks."""
+    for lock_fd in lock_fds:
+        try:
+            fcntl.flock(lock_fd, fcntl.LOCK_UN)
+            os.close(lock_fd)
+            logger.debug("Released initialization lock (fd=%s)", lock_fd)
+        except (OSError, ValueError):
+            pass
+
+
+def load_omni_transfer_config_for_model(model: str, config_path: str | None) -> Any:
+    """Load omni transfer config from an explicit path or resolved model config."""
+    from vllm_omni.distributed.omni_connectors import load_omni_transfer_config
+
+    try:
+        resolved_config_path = config_path or resolve_model_config_path(model)
+        return load_omni_transfer_config(resolved_config_path)
+    except Exception as e:
+        logger.warning("[stage_init] Failed to load transfer config: %s", e)
+        return None
+
+
+def get_stage_connector_spec(
+    omni_transfer_config: Any,
+    stage_id: int,
+    async_chunk: bool,
+) -> dict[str, Any]:
+    """Return the first connector spec for the stage when async chunking is enabled."""
+    from vllm_omni.distributed.omni_connectors import get_stage_connector_config
+
+    if not async_chunk:
+        return {}
+
+    stage_connectors_cfg = get_stage_connector_config(omni_transfer_config, stage_id)
+    for cfg in stage_connectors_cfg.values():
+        return dict(cfg.get("spec", {}))
+    return {}
+
+
+def initialize_diffusion_stage(model: str, stage_cfg: Any, metadata: StageMetadata) -> Any:
+    """Build a diffusion stage client."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+
+    od_config = OmniDiffusionConfig.from_kwargs(
+        model=model,
+        **_to_dict(stage_cfg.engine_args),
+    )
+    if metadata.cfg_kv_collect_func is not None:
+        od_config.cfg_kv_collect_func = metadata.cfg_kv_collect_func
+    return StageDiffusionClient(model, od_config, metadata)
+
+
+def close_started_llm_stage(started: StartedLlmStage) -> None:
+    """Close managers owned by a launched stage that never attached."""
+    resources = (
+        ("engine manager", started.engine_manager),
+        ("coordinator", started.coordinator),
+    )
+    for resource_name, resource in resources:
+        if resource is None:
+            continue
+        try:
+            resource.close()
+        except Exception as cleanup_error:
+            logger.warning(
+                "[stage_init] Failed to close launched %s for stage %s: %s",
+                resource_name,
+                started.stage_id,
+                cleanup_error,
+            )
+
+
+def finalize_initialized_stages(
+    stage_clients: list[Any | None],
+    input_processor: InputProcessor | None,
+) -> tuple[list[Any], list[Any], list[dict[str, Any]]]:
+    """Validate successful init and build runtime metadata lists."""
+    if any(stage_client is None for stage_client in stage_clients):
+        raise RuntimeError("Stage initialization completed with missing stage clients")
+
+    initialized_stage_clients = [stage_client for stage_client in stage_clients if stage_client is not None]
+    default_sampling_params_list = [stage_client.default_sampling_params for stage_client in initialized_stage_clients]
+    stage_metadata = [
+        {
+            "final_output": stage_client.final_output,
+            "final_output_type": stage_client.final_output_type,
+            "stage_type": stage_client.stage_type,
+        }
+        for stage_client in initialized_stage_clients
+    ]
+
+    if not isinstance(input_processor, InputProcessor):
+        has_llm_stage = any(metadata.get("stage_type") != "diffusion" for metadata in stage_metadata)
+        if has_llm_stage:
+            raise RuntimeError("Failed to initialize stage-0 InputProcessor for LLM pipeline")
+
+    return initialized_stage_clients, default_sampling_params_list, stage_metadata
+
+
+def cleanup_failed_stage_initialization(
+    stage_clients: list[Any | None],
+    started_llm_stages: list[StartedLlmStage],
+) -> None:
+    """Shutdown attached stages and close any launched-but-unattached engines."""
+    for cleanup_stage_id, stage_client in reversed(list(enumerate(stage_clients))):
+        if stage_client is None:
+            continue
+        try:
+            stage_client.shutdown()
+        except Exception as cleanup_error:
+            logger.warning(
+                "[stage_init] Failed to shutdown initialized stage %s after init failure: %s",
+                cleanup_stage_id,
+                cleanup_error,
+            )
+
+    for started in reversed(started_llm_stages):
+        if stage_clients[started.stage_id] is not None:
+            continue
+        close_started_llm_stage(started)

vllm_omni/entrypoints/__init__.py

@@ -3,11 +3,6 @@
 
 """
 vLLM-Omni entrypoints module.
-
-Provides high-level interfaces for running omni models including:
-- AsyncOmni: Async orchestrator for multi-stage LLM pipelines
-- AsyncOmniDiffusion: Async interface for diffusion model inference
-- Omni: Unified entrypoint that auto-selects between LLM and Diffusion
 """
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni

vllm_omni/entrypoints/async_omni_diffusion.py

@@ -10,6 +10,7 @@ enabling concurrent request handling and streaming generation.
 
 import asyncio
 import uuid
+import weakref
 from collections.abc import AsyncGenerator, Iterable
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any
@@ -27,6 +28,18 @@ from vllm_omni.outputs import OmniRequestOutput
 logger = init_logger(__name__)
 
 
+def _weak_close_async_omni_diffusion(engine: DiffusionEngine, executor: ThreadPoolExecutor) -> None:
+    """Best-effort diffusion cleanup for GC finalization."""
+    try:
+        engine.close()
+    except Exception:
+        pass
+    try:
+        executor.shutdown(wait=False)
+    except Exception:
+        pass
+
+
 class AsyncOmniDiffusion:
     """Async entry point for vLLM-Omni diffusion model inference.
 
@@ -127,6 +140,12 @@ class AsyncOmniDiffusion:
         # Thread pool for running sync engine in async context
         self._executor = ThreadPoolExecutor(max_workers=1)
         self._closed = False
+        self._weak_finalizer = weakref.finalize(
+            self,
+            _weak_close_async_omni_diffusion,
+            self.engine,
+            self._executor,
+        )
 
         logger.info("AsyncOmniDiffusion initialized with model: %s", model)
 
@@ -217,6 +236,9 @@ class AsyncOmniDiffusion:
         if self._closed:
             return
         self._closed = True
+        finalizer = getattr(self, "_weak_finalizer", None)
+        if finalizer is not None and finalizer.alive:
+            finalizer.detach()
 
         try:
             self.engine.close()
@@ -234,13 +256,6 @@ class AsyncOmniDiffusion:
         """Alias for close() method."""
         self.close()
 
-    def __del__(self) -> None:
-        """Best-effort cleanup on deletion."""
-        try:
-            self.close()
-        except Exception:
-            pass
-
     async def abort(self, request_id: str | Iterable[str]) -> None:
         """Abort a request."""
         self.engine.abort(request_id)

vllm_omni/entrypoints/async_omni_llm.py

@@ -1,216 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import asyncio
-import os
-import socket
-from typing import TYPE_CHECKING
-
-import torch
-from vllm.config import VllmConfig
-from vllm.logger import init_logger
-from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
-from vllm.renderers import renderer_from_config
-from vllm.tracing import init_tracer
-from vllm.transformers_utils.config import maybe_register_config_serialize_by_value
-from vllm.usage.usage_lib import UsageContext
-from vllm.utils.func_utils import deprecate_kwargs
-from vllm.v1.engine.async_llm import AsyncLLM
-from vllm.v1.engine.core_client import EngineCoreClient
-from vllm.v1.executor.abstract import Executor
-from vllm.v1.metrics.loggers import StatLoggerFactory, StatLoggerManager
-
-from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs
-from vllm_omni.engine.input_processor import OmniInputProcessor
-from vllm_omni.engine.output_processor import MultimodalOutputProcessor
-
-if TYPE_CHECKING:
-    pass
-
-logger = init_logger(__name__)
-
-
-class AsyncOmniLLM(AsyncLLM):
-    """Async single-stage LLM engine for use within a stage worker process.
-
-    This class extends the base vLLM AsyncLLM class with omni-specific
-    processors for handling multimodal inputs and outputs. It is used
-    internally by AsyncOmniStage workers and should not be instantiated
-    directly by users.
-
-    Args:
-        engine_args: AsyncOmniEngineArgs containing engine configuration
-        vllm_config: Global vLLM configuration
-        executor_class: Executor implementation class, e.g. MultiprocExecutor
-        log_stats: Whether to log statistics
-        usage_context: Usage context of the LLM (default: ENGINE_CONTEXT)
-        mm_registry: Multi-modal registry for processing multimodal inputs
-        use_cached_outputs: Whether to use cached outputs
-        log_requests: Whether to log requests
-        start_engine_loop: Whether to start the engine loop automatically
-        stat_loggers: Customized stat loggers for the engine.
-            If not provided, default stat loggers will be used.
-            Note: Stat logger interface may change in V1.
-        client_addresses: Optional dictionary mapping client names to addresses
-        client_count: Total number of clients (default: 1)
-        client_index: Index of this client (default: 0)
-    """
-
-    def __init__(
-        self,
-        engine_args: AsyncOmniEngineArgs,
-        vllm_config: VllmConfig,
-        executor_class: type[Executor],
-        log_stats: bool,
-        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
-        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
-        use_cached_outputs: bool = False,
-        log_requests: bool = True,
-        start_engine_loop: bool = True,
-        stat_loggers: list[StatLoggerFactory] | None = None,
-        client_addresses: dict[str, str] | None = None,
-        client_count: int = 1,
-        client_index: int = 0,
-    ) -> None:
-        """
-        Create an AsyncOmniLLM.
-
-        Args:
-            vllm_config: global configuration.
-            executor_class: an Executor impl, e.g. MultiprocExecutor.
-            log_stats: Whether to log stats.
-            usage_context: Usage context of the LLM.
-            mm_registry: Multi-modal registry.
-            use_cached_outputs: Whether to use cached outputs.
-            log_requests: Whether to log requests.
-            start_engine_loop: Whether to start the engine loop.
-            stat_loggers: customized stat loggers for the engine.
-                If not provided, default stat loggers will be used.
-                PLEASE BE AWARE THAT STAT LOGGER IS NOT STABLE
-                IN V1, AND ITS BASE CLASS INTERFACE MIGHT CHANGE.
-
-        Returns:
-            None
-        """
-        # Ensure we can serialize custom transformer configs
-        maybe_register_config_serialize_by_value()
-
-        self.model_config = vllm_config.model_config
-        self.vllm_config = vllm_config
-        self.observability_config = vllm_config.observability_config
-        self.log_requests = log_requests
-
-        self.log_stats = log_stats or (stat_loggers is not None)
-        if not log_stats and stat_loggers is not None:
-            logger.info(
-                "AsyncLLM created with log_stats=False and non-empty custom logger list; "
-                "enabling logging without default stat loggers"
-            )
-
-        # InputProcessor (converts Inputs --> EngineCoreRequests).
-        self.input_processor = OmniInputProcessor(
-            vllm_config=vllm_config,
-            mm_registry=mm_registry,
-        )
-
-        self.renderer = renderer_from_config(self.vllm_config)
-
-        # OutputProcessor (converts EngineCoreOutputs --> RequestOutput).
-        self.output_processor = MultimodalOutputProcessor(
-            tokenizer=self.renderer.tokenizer,
-            log_stats=self.log_stats,
-            engine_core_output_type=engine_args.engine_output_type,
-        )
-
-        if self.observability_config.otlp_traces_endpoint is not None:
-            tracer = init_tracer("vllm.llm_engine", self.observability_config.otlp_traces_endpoint)
-            self.output_processor.tracer = tracer
-
-        # Pause / resume state for async RL workflows.
-        self._pause_cond = asyncio.Condition()
-        self._paused = False
-
-        # EngineCore (starts the engine in background process).
-        self.engine_core = EngineCoreClient.make_async_mp_client(
-            vllm_config=vllm_config,
-            executor_class=executor_class,
-            log_stats=self.log_stats,
-            client_addresses=client_addresses,
-            client_count=client_count,
-            client_index=client_index,
-        )
-
-        # Loggers.
-        self.logger_manager: StatLoggerManager | None = None
-        if self.log_stats:
-            self.logger_manager = StatLoggerManager(
-                vllm_config=vllm_config,
-                engine_idxs=self.engine_core.engine_ranks_managed,
-                custom_stat_loggers=stat_loggers,
-                enable_default_loggers=log_stats,
-                client_count=client_count,
-            )
-            self.logger_manager.log_engine_initialized()
-
-        self.output_handler: asyncio.Task | None = None
-        try:
-            # Start output handler eagerly if we are in the asyncio eventloop.
-            asyncio.get_running_loop()
-            self._run_output_handler()
-        except RuntimeError:
-            pass
-
-        # Use profiler_config from vllm_config (new way, aligned with vllm v1)
-        if vllm_config.profiler_config.profiler == "torch" and not vllm_config.profiler_config.ignore_frontend:
-            profiler_dir = vllm_config.profiler_config.torch_profiler_dir
-            logger.info(
-                "Torch profiler enabled. AsyncOmniLLM CPU traces will be collected under %s",
-                profiler_dir,
-            )
-            worker_name = f"{socket.gethostname()}_{os.getpid()}.async_omni_llm"
-            self.profiler = torch.profiler.profile(
-                activities=[
-                    torch.profiler.ProfilerActivity.CPU,
-                ],
-                with_stack=vllm_config.profiler_config.torch_profiler_with_stack,
-                on_trace_ready=torch.profiler.tensorboard_trace_handler(
-                    profiler_dir,
-                    worker_name=worker_name,
-                    use_gzip=vllm_config.profiler_config.torch_profiler_use_gzip,
-                ),
-            )
-        else:
-            self.profiler = None
-
-    @classmethod
-    @deprecate_kwargs(
-        "disable_log_requests",
-        additional_message=("This argument will have no effect. Use `enable_log_requests` instead."),
-    )
-    def from_vllm_config(
-        cls,
-        vllm_config: VllmConfig,
-        engine_args: AsyncOmniEngineArgs,
-        start_engine_loop: bool = True,
-        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
-        stat_loggers: list[StatLoggerFactory] | None = None,
-        enable_log_requests: bool = False,
-        disable_log_stats: bool = False,
-        client_addresses: dict[str, str] | None = None,
-        client_count: int = 1,
-        client_index: int = 0,
-        disable_log_requests: bool = True,  # Deprecated, will be removed
-    ) -> "AsyncLLM":
-        # Create the LLMEngine.
-        return cls(
-            vllm_config=vllm_config,
-            executor_class=Executor.get_class(vllm_config),
-            start_engine_loop=start_engine_loop,
-            stat_loggers=stat_loggers,
-            log_requests=enable_log_requests,
-            log_stats=not disable_log_stats,
-            usage_context=usage_context,
-            client_addresses=client_addresses,
-            client_count=client_count,
-            client_index=client_index,
-            engine_args=engine_args,
-        )

vllm_omni/entrypoints/cli/serve.py

@@ -8,37 +8,20 @@ diffusion models (e.g., Qwen-Image) through the same CLI interface.
 import argparse
 import json
 import os
-import signal
 from typing import Any
 
-import msgspec.msgpack
 import uvloop
-import zmq
 from vllm.entrypoints.cli.types import CLISubcommand
 from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
 from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
 from vllm.logger import init_logger
 from vllm.utils.argparse_utils import FlexibleArgumentParser
-from vllm.utils.network_utils import make_zmq_socket
-from vllm.v1.utils import get_engine_client_zmq_addr
 
-from vllm_omni.distributed.omni_connectors import (
-    get_connectors_config_for_stage,
-    load_omni_transfer_config,
-)
-from vllm_omni.distributed.omni_connectors.utils.initialization import (
-    resolve_omni_kv_config_for_stage,
-)
 from vllm_omni.entrypoints.cli.logo import log_logo
-from vllm_omni.entrypoints.omni import OmniBase, omni_snapshot_download
-from vllm_omni.entrypoints.omni_stage import OmniStage
 from vllm_omni.entrypoints.openai.api_server import omni_run_server
-from vllm_omni.entrypoints.utils import inject_omni_kv_config
 
 logger = init_logger(__name__)
 
-HANDSHAKE_TIMEOUT_MINS = 5
-
 DESCRIPTION = """Launch a local OpenAI-compatible API server to serve Omni models
 via HTTP. Supports both multi-stage LLM models and diffusion models.
 
@@ -372,123 +355,33 @@ class OmniServeCommand(CLISubcommand):
 
 
 def _create_default_diffusion_stage_cfg(args: argparse.Namespace) -> list[dict[str, Any]]:
-    omni_base = OmniBase.__new__(OmniBase)
-    return omni_base._create_default_diffusion_stage_cfg(vars(args))
+    """Create default diffusion stage configuration.
+
+    Uses AsyncOmniEngine's implementation which doesn't have OmegaConf
+    compatibility issues.
+    """
+    from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+
+    return AsyncOmniEngine._create_default_diffusion_stage_cfg(vars(args))
 
 
 def run_headless(args: argparse.Namespace) -> None:
-    if args.api_server_count is not None and args.api_server_count > 1:
-        raise ValueError("api_server_count can't be set in headless mode")
-    if args.worker_backend != "multi_process":
-        raise ValueError("headless mode requires worker_backend=multi_process")
+    """Run a single stage in headless mode.
 
-    model = omni_snapshot_download(args.model)
+    .. deprecated:: 0.x.x
+        Headless mode is deprecated and will be removed in a future version.
+        It is only compatible with the old OmniStage-based runtime.
+        The current AsyncOmniEngine-based runtime does not support headless mode.
 
-    omni_base = OmniBase.__new__(OmniBase)
-    args_dict = vars(args).copy()
-    args_dict["model"] = model
-    config_path, stage_configs = omni_base._resolve_stage_configs(model, args_dict)
-
-    single_stage_id = args.stage_id
-    if single_stage_id is None:
-        if len(stage_configs) != 1:
-            raise ValueError("--stage-id is required in headless mode for multi-stage configs")
-        single_stage_id = getattr(stage_configs[0], "stage_id", 0)
-
-    stage_config = None
-    for cfg in stage_configs:
-        if getattr(cfg, "stage_id", None) == single_stage_id:
-            stage_config = cfg
-            break
-    if stage_config is None:
-        raise ValueError(f"No stage matches stage_id={single_stage_id}.")
-
-    # TODO(wuhang): Support connectors config by cli
-    transfer_config = load_omni_transfer_config(config_path, default_shm_threshold=args.shm_threshold_bytes)
-    connectors_config = get_connectors_config_for_stage(transfer_config, single_stage_id)
-
-    omni_master_address = args.omni_master_address
-    omni_master_port = args.omni_master_port
-
-    # Perform handshake with orchestrator to get dynamically allocated endpoints
-    with zmq.Context() as zmq_ctx:
-        handshake_endpoint = get_engine_client_zmq_addr(
-            local_only=False, host=omni_master_address, port=omni_master_port
-        )
-
-        with make_zmq_socket(zmq_ctx, handshake_endpoint, zmq.REQ, bind=False, linger=5000) as handshake_socket:
-            # TODO(wuhang): Define protocol in python dataclass.
-            handshake_msg = {"type": "handshake", "stage_id": single_stage_id}
-            handshake_socket.send(msgspec.msgpack.encode(handshake_msg))
-
-            # Wait for response with timeout
-            if not handshake_socket.poll(timeout=HANDSHAKE_TIMEOUT_MINS * 60_000):
-                raise RuntimeError(
-                    f"Handshake timeout ({HANDSHAKE_TIMEOUT_MINS} minutes) for stage-{single_stage_id} "
-                    f"at {handshake_endpoint}"
-                )
-
-            try:
-                response = msgspec.msgpack.decode(handshake_socket.recv())
-            except msgspec.DecodeError as exc:
-                raise RuntimeError(
-                    f"Handshake decode failed for stage-{single_stage_id} at {handshake_endpoint}: {exc}"
-                ) from exc
-            except Exception as exc:  # pragma: no cover - unexpected decode errors
-                raise RuntimeError(
-                    f"Unexpected error decoding handshake for stage-{single_stage_id} at {handshake_endpoint}: {exc}"
-                ) from exc
-
-            if not response["ok"]:
-                error_msg = response["error"]
-                raise RuntimeError(f"Handshake failed for stage-{single_stage_id}: {error_msg}")
-
-            in_endpoint, out_endpoint = response["in_endpoint"], response["out_endpoint"]
-
-            logger.info(
-                f"[Headless] Stage-{single_stage_id} received endpoints via handshake: "
-                f"in={in_endpoint}, out={out_endpoint}"
-            )
-
-    shutdown_requested = False
-
-    def signal_handler(signum, frame):
-        nonlocal shutdown_requested
-        if shutdown_requested:
-            return
-        shutdown_requested = True
-        raise SystemExit
-
-    signal.signal(signal.SIGTERM, signal_handler)
-    signal.signal(signal.SIGINT, signal_handler)
-
-    stage = OmniStage(stage_config, stage_init_timeout=args.stage_init_timeout)
-    stage.attach_queues(in_endpoint, out_endpoint)
-
-    # Inject YAML-resolved connector config into omni_kv_config for in-engine usage.
-    try:
-        omni_conn_cfg, omni_from, omni_to = resolve_omni_kv_config_for_stage(transfer_config, single_stage_id)
-        if omni_conn_cfg:
-            inject_omni_kv_config(stage, omni_conn_cfg, omni_from, omni_to)  # type: ignore
-    except Exception as e:
-        logger.debug("[Headless] Failed to inject omni connector config into stage-%s: %s", single_stage_id, e)
-
-    old_env = os.environ.get("VLLM_LOGGING_PREFIX")
-    os.environ["VLLM_LOGGING_PREFIX"] = f"[Stage-{single_stage_id}] {'' if old_env is None else old_env}"
-    try:
-        stage.init_stage_worker(
-            model,
-            is_async=True,
-            shm_threshold_bytes=int(args.shm_threshold_bytes),
-            batch_timeout=int(args.batch_timeout),
-            connectors_config=connectors_config,
-            worker_backend="multi_process",
-            ignore_runtime_config=True,
-        )
-        if stage._proc is not None:
-            stage._proc.join()
-    finally:
-        stage.stop_stage_worker()
+    Raises:
+        RuntimeError: Always raises an error indicating headless mode is deprecated.
+    """
+    raise RuntimeError(
+        "Headless mode is deprecated and not supported in the current runtime. "
+        "Please use the standard orchestrator mode (without --headless flag). "
+        "If you need distributed deployment, consider using Ray backend or "
+        "other distributed serving solutions."
+    )
 
 
 def cmd_init() -> list[CLISubcommand]:

vllm_omni/entrypoints/omni_base.py

@@ -0,0 +1,290 @@
+from __future__ import annotations
+
+import os
+import time
+import types
+import weakref
+from collections.abc import Sequence
+from pprint import pformat
+from typing import TYPE_CHECKING, Any, Literal
+
+import huggingface_hub
+from vllm.logger import init_logger
+from vllm.v1.engine.exceptions import EngineDeadError
+
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.entrypoints.client_request_state import ClientRequestState
+from vllm_omni.entrypoints.utils import get_final_stage_id_for_e2e
+from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics
+from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
+from vllm_omni.outputs import OmniRequestOutput
+
+if TYPE_CHECKING:
+    from vllm_omni.engine.arg_utils import OmniEngineArgs
+
+logger = init_logger(__name__)
+
+
+def _weak_shutdown_engine(engine: AsyncOmniEngine) -> None:
+    """Best-effort engine cleanup for GC finalization."""
+    try:
+        engine.shutdown()
+    except Exception:
+        pass
+
+
+def omni_snapshot_download(model_id: str) -> str:
+    if os.path.exists(model_id):
+        return model_id
+
+    # TODO: this is just a workaround for quickly use modelscope, we should support
+    # modelscope in weight loading feature instead of using `snapshot_download`
+    if os.environ.get("VLLM_USE_MODELSCOPE", False):
+        from modelscope.hub.snapshot_download import snapshot_download
+
+        return snapshot_download(model_id)
+
+    try:
+        download_weights_from_hf_specific(
+            model_name_or_path=model_id,
+            cache_dir=None,
+            allow_patterns=["*"],
+            require_all=True,
+        )
+    except huggingface_hub.errors.RepositoryNotFoundError:
+        logger.warning("Repository not found for '%s'.", model_id)
+
+    return model_id
+
+
+OutputMessageHandleResult = tuple[Literal[True], None, None, None] | tuple[Literal[False], str, int, ClientRequestState]
+
+
+class OmniBase:
+    """Shared runtime foundation for AsyncOmni and Omni."""
+
+    def __init__(
+        self,
+        model: str,
+        **kwargs: Any,
+    ) -> None:
+        engine_args: OmniEngineArgs | None = kwargs.pop("engine_args", None)
+        stage_init_timeout = kwargs.pop("stage_init_timeout", 300)
+        init_timeout = kwargs.pop("init_timeout", 600)
+        log_stats = kwargs.pop("log_stats", False)
+        async_chunk = kwargs.pop("async_chunk", False)
+        output_modalities = kwargs.pop("output_modalities", None)
+
+        if "log_requests" in kwargs:
+            raise TypeError("`log_requests` has been removed in Omni/AsyncOmni. Use `log_stats`.")
+        model = omni_snapshot_download(model)
+        self.model = model
+        self.log_stats = log_stats
+        self.async_chunk = async_chunk
+        self.output_modalities = output_modalities or []
+
+        logger.info("[%s] Initializing with model %s", self.__class__.__name__, model)
+        st = time.time()
+        self.engine = AsyncOmniEngine(
+            model=model,
+            engine_args=engine_args,
+            init_timeout=init_timeout,
+            stage_init_timeout=stage_init_timeout,
+            **kwargs,
+        )
+        self._shutdown_called = False
+        self._weak_finalizer = weakref.finalize(self, _weak_shutdown_engine, self.engine)
+        et = time.time()
+        logger.info("[%s] AsyncOmniEngine initialized in %.2f seconds", self.__class__.__name__, et - st)
+        self.async_chunk = bool(self.async_chunk or getattr(self.engine, "async_chunk", False))
+
+        self.request_states: dict[str, ClientRequestState] = {}
+
+        self.default_sampling_params_list = self.engine.default_sampling_params_list
+        if not self.output_modalities:
+            self.output_modalities = [
+                self.engine.get_stage_metadata(i).get("final_output_type") for i in range(self.engine.num_stages)
+            ]
+
+        self._stage_meta_list = [
+            types.SimpleNamespace(**self.engine.get_stage_metadata(i)) for i in range(self.engine.num_stages)
+        ]
+
+        logger.info(
+            "[%s] Initialized with %s stages for model %s",
+            self.__class__.__name__,
+            self.engine.num_stages,
+            model,
+        )
+
+    @property
+    def num_stages(self) -> int:
+        return self.engine.num_stages
+
+    @property
+    def is_running(self) -> bool:
+        return self.engine.is_alive()
+
+    def check_health(self) -> None:
+        if not self.engine.is_alive():
+            raise EngineDeadError("Orchestrator process is not alive")
+
+    def resolve_sampling_params_list(
+        self,
+        sampling_params_list: Sequence[Any] | Any | None,
+    ) -> Sequence[Any]:
+        if sampling_params_list is None:
+            normalized = self.default_sampling_params_list
+        elif isinstance(sampling_params_list, Sequence) and not isinstance(sampling_params_list, (str, bytes)):
+            normalized = sampling_params_list
+        elif self.num_stages == 1:
+            normalized = [sampling_params_list]
+        else:
+            raise ValueError(f"Expected {self.num_stages} sampling params, got a single sampling params object")
+        if len(normalized) != self.num_stages:
+            raise ValueError(f"Expected {self.num_stages} sampling params, got {len(normalized)}")
+        return normalized
+
+    def _log_summary_and_cleanup(self, request_id: str) -> None:
+        req_state = self.request_states.get(request_id)
+        try:
+            if req_state is None or req_state.metrics is None:
+                return
+            summary = req_state.metrics.build_and_log_summary()
+            logger.info("[Summary] %s", pformat(summary, sort_dicts=False))
+        except Exception:
+            logger.exception(
+                "[%s] Failed to build/log summary for req=%s",
+                self.__class__.__name__,
+                request_id,
+            )
+        finally:
+            self.request_states.pop(request_id, None)
+
+    def _compute_final_stage_id(self, output_modalities: list[str] | None) -> int:
+        return get_final_stage_id_for_e2e(
+            output_modalities,
+            self.output_modalities,
+            self._stage_meta_list,
+        )
+
+    def _process_stage_metrics_message(self, msg: dict[str, Any]) -> None:
+        req_id = msg.get("request_id")
+        req_state = self.request_states.get(req_id)
+        if req_state is None or req_state.metrics is None:
+            return
+        _m = msg.get("metrics")
+        if _m is None:
+            return
+        stage_id = msg.get("stage_id", 0)
+        req_state.metrics.on_stage_metrics(stage_id, req_id, _m)
+        submit_ts = msg.get("stage_submit_ts")
+        now = time.time()
+        if req_state.metrics.stage_first_ts[stage_id] is None:
+            req_state.metrics.stage_first_ts[stage_id] = submit_ts if submit_ts is not None else now
+        req_state.metrics.stage_last_ts[stage_id] = max(req_state.metrics.stage_last_ts[stage_id] or 0.0, now)
+
+    def _handle_output_message(
+        self,
+        msg: dict[str, Any] | None,
+    ) -> OutputMessageHandleResult:
+        """Handle one Orchestrator output-queue message."""
+        if msg is None:
+            return True, None, None, None
+
+        msg_type = msg.get("type")
+        if msg_type == "stage_metrics":
+            self._process_stage_metrics_message(msg)
+            return True, None, None, None
+
+        if msg_type == "error":
+            raise RuntimeError(msg.get("error", "Orchestrator returned an error message"))
+
+        if msg_type != "output":
+            logger.warning("[%s] got unexpected msg type: %s", self.__class__.__name__, msg_type)
+            return True, None, None, None
+
+        req_id = msg.get("request_id")
+        if req_id is None:
+            logger.warning("[%s] got output message without request_id", self.__class__.__name__)
+            return True, None, None, None
+
+        stage_id = msg.get("stage_id")
+        if stage_id is None:
+            logger.warning("[%s] got output message without stage_id for req=%s", self.__class__.__name__, req_id)
+            return True, None, None, None
+
+        req_state = self.request_states.get(req_id)
+        if req_state is None:
+            logger.debug(
+                "[%s] dropping output for unknown req %s",
+                self.__class__.__name__,
+                req_id,
+            )
+            return True, None, None, None
+
+        req_state.stage_id = stage_id
+
+        return False, req_id, stage_id, req_state
+
+    def _process_single_result(
+        self,
+        result: dict[str, Any],
+        stage_id: int,
+        metrics: OrchestratorMetrics,
+        req_start_ts: dict[str, float],
+        wall_start_ts: float,
+        final_stage_id_for_e2e: int,
+    ) -> OmniRequestOutput | None:
+        req_id = result.get("request_id")
+        engine_outputs = result.get("engine_outputs")
+        finished = engine_outputs.finished
+
+        submit_ts = result.get("stage_submit_ts")
+        now = time.time()
+        if metrics.stage_first_ts[stage_id] is None:
+            metrics.stage_first_ts[stage_id] = submit_ts if submit_ts is not None else now
+        metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, now)
+
+        _m = result.get("metrics")
+        if finished and _m is not None:
+            metrics.on_stage_metrics(stage_id, req_id, _m)
+
+        stage_meta = self.engine.get_stage_metadata(stage_id)
+        if not stage_meta["final_output"]:
+            return None
+
+        try:
+            rid_key = str(req_id)
+            if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done and finished:
+                metrics.on_finalize_request(
+                    stage_id,
+                    req_id,
+                    req_start_ts.get(req_id, wall_start_ts),
+                )
+        except Exception:
+            logger.exception("[%s] Finalize request handling error", self.__class__.__name__)
+
+        images = getattr(engine_outputs, "images", []) if stage_meta["final_output_type"] == "image" else []
+        return OmniRequestOutput(
+            stage_id=stage_id,
+            final_output_type=stage_meta["final_output_type"],
+            request_output=engine_outputs,
+            images=images,
+        )
+
+    def shutdown(self) -> None:
+        logger.info("[%s] Shutting down", self.__class__.__name__)
+        self._shutdown_base()
+
+    def close(self) -> None:
+        self.shutdown()
+
+    def _shutdown_base(self) -> None:
+        if getattr(self, "_shutdown_called", False):
+            return
+        self._shutdown_called = True
+        finalizer = getattr(self, "_weak_finalizer", None)
+        if finalizer is not None and finalizer.alive:
+            finalizer.detach()
+        self.engine.shutdown()

vllm_omni/entrypoints/omni_diffusion.py

@@ -1,169 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import time
-import uuid
-from collections.abc import Sequence
-
-from vllm.logger import init_logger
-from vllm.transformers_utils.config import get_hf_file_to_dict
-
-from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
-from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
-from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
-from vllm_omni.outputs import OmniRequestOutput
-
-logger = init_logger(__name__)
-
-
-class OmniDiffusion:
-    """
-    It is the main class to interact with vLLM-Omni diffusion models.
-    It acts as a high-level interface that prepares requests and
-    delegates the actual diffusion process to the DiffusionEngine.
-
-    You can pass either an `OmniDiffusionConfig` via `od_config`, or
-    pass kwargs such as `model="Qwen/Qwen-Image"`,
-    which will be forwarded to `OmniDiffusionConfig.from_kwargs`.
-    """
-
-    def __init__(self, od_config: OmniDiffusionConfig | None = None, **kwargs):
-        # Capture stage info from kwargs before they might be filtered out
-        stage_id = kwargs.get("stage_id")
-        engine_input_source = kwargs.get("engine_input_source")
-        cfg_kv_collect_func = kwargs.pop("cfg_kv_collect_func", None)
-
-        if od_config is None:
-            od_config = OmniDiffusionConfig.from_kwargs(**kwargs)
-        elif isinstance(od_config, dict):
-            # If config is dict, check it too (priority to kwargs if both exist)
-            if stage_id is None:
-                stage_id = od_config.get("stage_id")
-            if engine_input_source is None:
-                engine_input_source = od_config.get("engine_input_source")
-            od_config = OmniDiffusionConfig.from_kwargs(**od_config)
-
-        self.od_config = od_config
-
-        # Inject stage info into omni_kv_config if present
-        if stage_id is not None:
-            self.od_config.omni_kv_config.setdefault("stage_id", stage_id)
-        if engine_input_source is not None:
-            self.od_config.omni_kv_config.setdefault("engine_input_source", engine_input_source)
-
-        # Detect model class and load config
-        # Diffusers-style models expose `model_index.json` with `_class_name`.
-        # Non-diffusers models (e.g. Bagel, NextStep, GLM-Image) only have `config.json`,
-        # so we fall back to reading that and mapping model_type manually.
-        try:
-            config_dict = get_hf_file_to_dict(
-                "model_index.json",
-                od_config.model,
-            )
-
-            if config_dict is None:
-                raise FileNotFoundError("model_index.json not found")
-
-            if od_config.model_class_name is None:
-                od_config.model_class_name = config_dict.get("_class_name", None)
-            od_config.update_multimodal_support()
-
-            if od_config.model_class_name == "DreamIDOmniPipeline":
-                od_config.model_config = config_dict
-            else:
-                tf_config_dict = get_hf_file_to_dict(
-                    "transformer/config.json",
-                    od_config.model,
-                )
-                od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict)
-
-        except (AttributeError, OSError, ValueError, FileNotFoundError):
-            cfg = get_hf_file_to_dict("config.json", od_config.model)
-            if cfg is None:
-                raise ValueError(f"Could not find config.json or model_index.json for model {od_config.model}")
-
-            # Map model_type or architecture to pipeline class
-            model_type = cfg.get("model_type")
-            architectures = cfg.get("architectures") or []
-            pipeline_class = None
-            # Bagel/NextStep models don't have a model_index.json, so we set the pipeline class name manually
-            if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
-                pipeline_class = "BagelPipeline"
-            elif model_type == "nextstep":
-                if od_config.model_class_name is None:
-                    pipeline_class = "NextStep11Pipeline"
-            elif model_type == "glm-image" or "GlmImageForConditionalGeneration" in architectures:
-                pipeline_class = "GlmImagePipeline"
-            elif architectures and len(architectures) == 1:
-                pipeline_class = architectures[0]
-
-            if pipeline_class is None:
-                raise ValueError(f"Unknown model type: {model_type}, architectures: {architectures}")
-
-            if od_config.model_class_name is None:
-                od_config.model_class_name = pipeline_class
-            od_config.tf_model_config = TransformerConfig().from_dict(cfg)
-            od_config.update_multimodal_support()
-
-        if cfg_kv_collect_func is not None:
-            od_config.cfg_kv_collect_func = cfg_kv_collect_func
-
-        self.engine: DiffusionEngine = DiffusionEngine.make_engine(od_config)
-
-    def generate(
-        self,
-        prompts: OmniPromptType | Sequence[OmniPromptType],
-        sampling_params: OmniDiffusionSamplingParams,
-        request_ids: list[str] = [],
-    ) -> list[OmniRequestOutput]:
-        _t0 = time.perf_counter()
-        if isinstance(prompts, (str, dict)):
-            prompts = [prompts]
-        else:
-            prompts = list(prompts)
-
-        # Check if request_id is provided in kwargs
-        if len(request_ids) < len(prompts):
-            request_ids.extend(f"{i + len(request_ids)}_{uuid.uuid4()}" for i in range(len(prompts) - len(request_ids)))
-
-        request = OmniDiffusionRequest(prompts, sampling_params, request_ids)
-        result = self._run_engine(request)
-        _t_ms = (time.perf_counter() - _t0) * 1000
-        logger.info("OmniDiffusion.generate total: %.2f ms", _t_ms)
-        return result
-
-    def _run_engine(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
-        return self.engine.step(request)
-
-    def close(self) -> None:
-        self.engine.close()
-
-    def __del__(self):  # pragma: no cover - best effort cleanup
-        try:
-            self.close()
-        except Exception:
-            pass
-
-    def start_profile(self, trace_filename: str | None = None) -> None:
-        """Start profiling for the diffusion model.
-
-        Args:
-            trace_filename: Optional base filename for trace files.
-                           If None, a timestamp-based name will be generated.
-        """
-        if hasattr(self, "engine") and self.engine:
-            self.engine.start_profile(trace_filename)
-        else:
-            raise RuntimeError("Diffusion engine not initialized")
-
-    def stop_profile(self) -> dict:
-        """Stop profiling and return profiling results.
-
-        Returns:
-            Dictionary containing paths to trace and table files.
-        """
-        if hasattr(self, "engine") and self.engine:
-            return self.engine.stop_profile()
-        else:
-            raise RuntimeError("Diffusion engine not initialized")

vllm_omni/entrypoints/omni_llm.py

@@ -1,275 +0,0 @@
-from collections.abc import Callable, Sequence
-from typing import Any
-
-import cloudpickle
-from pydantic import ValidationError
-from tqdm import tqdm
-
-# External library imports (vLLM)
-from vllm.config import CompilationConfig, StructuredOutputsConfig, is_init_field
-from vllm.entrypoints.llm import LLM
-from vllm.inputs import ProcessorInputs, PromptType
-from vllm.logger import init_logger
-from vllm.outputs import PoolingRequestOutput, RequestOutput
-from vllm.plugins.io_processors import get_io_processor
-from vllm.renderers.inputs.preprocess import parse_model_prompt
-from vllm.usage.usage_lib import UsageContext
-from vllm.utils.counter import Counter
-from vllm.v1.engine.llm_engine import LLMEngine
-
-from vllm_omni.distributed.omni_connectors import initialize_orchestrator_connectors
-
-# Internal imports (our code)
-from vllm_omni.engine.arg_utils import OmniEngineArgs
-from vllm_omni.engine.input_processor import OmniInputProcessor, reinject_omni_fields
-from vllm_omni.engine.output_processor import MultimodalOutputProcessor
-from vllm_omni.entrypoints.utils import (
-    filter_dataclass_kwargs,
-    resolve_model_config_path,
-)
-
-logger = init_logger(__name__)
-
-
-class OmniLLM(LLM):
-    """Main entry point for vLLM-Omni inference.
-
-    This class extends the base vLLM LLM class with omni-specific
-    processors for handling multimodal inputs and outputs. It provides
-    configuration loading for multi-stage pipelines, while stage management
-    is handled by the Omni class.
-
-    Args:
-        model: Model name or path to load
-        log_stats: Whether to enable statistics logging
-        compilation_config: Optional compilation configuration. Can be an
-            integer (compilation level), dict, or CompilationConfig instance.
-        hf_overrides: Optional HuggingFace model configuration overrides
-        structured_outputs_config: Optional structured outputs configuration.
-            Can be a dict or StructuredOutputsConfig instance.
-        init_sleep_seconds: Number of seconds to sleep between starting
-            each stage process during initialization (used by Omni class)
-        shm_threshold_bytes: Threshold in bytes for using shared memory
-            for IPC. Objects larger than this threshold will use shared memory.
-        batch_timeout: Timeout in seconds for batching requests within a stage
-        init_timeout: Timeout in seconds for waiting for all stages to initialize
-        **kwargs: Additional keyword arguments passed to the base LLM class
-            and engine
-
-    Example:
-        >>> llm = OmniLLM(model="Qwen/Qwen2.5-Omni-7B")
-        >>> # Stage management is handled by Omni class
-    """
-
-    def __init__(
-        self,
-        model: str,
-        log_stats: bool = False,
-        compilation_config: int | dict[str, Any] | CompilationConfig | None = None,
-        hf_overrides: dict[str, Any] | None = None,
-        structured_outputs_config: dict[str, Any] | StructuredOutputsConfig | None = None,
-        init_sleep_seconds: int = 20,
-        shm_threshold_bytes: int = 65536,
-        batch_timeout: int = 10,
-        init_timeout: int = 300,
-        **kwargs: Any,
-    ):
-        """LLM constructor with omni-specific configuration loading."""
-        # Store stage management parameters (used by Omni class)
-        self.worker_backend = kwargs.get("worker_backend", "multi_process")
-        self.ray_address = kwargs.get("ray_address", None)
-        self.batch_timeout = batch_timeout
-        self.log_stats: bool = bool(log_stats)
-
-        # Resolve model config path for connectors
-        self.config_path = resolve_model_config_path(model)
-
-        # Initialize connectors
-        self.omni_transfer_config, self.connectors = initialize_orchestrator_connectors(
-            self.config_path, worker_backend=self.worker_backend, shm_threshold_bytes=shm_threshold_bytes
-        )
-
-        # Initialize LLM engine
-        if "disable_log_stats" not in kwargs:
-            kwargs["disable_log_stats"] = True
-
-        if "worker_cls" in kwargs:
-            worker_cls = kwargs["worker_cls"]
-            # if the worker_cls is not qualified string name,
-            # we serialize it using cloudpickle to avoid pickling issues
-            if isinstance(worker_cls, type):
-                kwargs["worker_cls"] = cloudpickle.dumps(worker_cls)
-
-        if "kv_transfer_config" in kwargs and isinstance(kwargs["kv_transfer_config"], dict):
-            from vllm.config.kv_transfer import KVTransferConfig
-
-            raw_config_dict = kwargs["kv_transfer_config"]
-            try:
-                kwargs["kv_transfer_config"] = KVTransferConfig(**raw_config_dict)
-            except ValidationError as e:
-                logger.error(
-                    "Failed to convert 'kv_transfer_config' dict to KVTransferConfig object. Dict: %s. Error: %s",
-                    raw_config_dict,
-                    e,
-                )
-                raise ValueError(f"Invalid 'kv_transfer_config' provided: {e}") from e
-
-        # Extract omni_kv_config from kwargs if present (injected by Omni)
-        omni_kv_config = kwargs.pop("omni_kv_config", None)
-
-        if compilation_config is not None:
-            if isinstance(compilation_config, int):
-                compilation_config_instance = CompilationConfig(level=compilation_config)
-            elif isinstance(compilation_config, dict):
-                compilation_config_instance = CompilationConfig(
-                    **{k: v for k, v in compilation_config.items() if is_init_field(CompilationConfig, k)}
-                )
-            else:
-                compilation_config_instance = compilation_config
-        else:
-            compilation_config_instance = CompilationConfig()
-
-        if structured_outputs_config is not None:
-            if isinstance(structured_outputs_config, dict):
-                structured_outputs_instance = StructuredOutputsConfig(
-                    **{k: v for k, v in structured_outputs_config.items() if is_init_field(StructuredOutputsConfig, k)}
-                )
-            else:
-                structured_outputs_instance = structured_outputs_config
-        else:
-            structured_outputs_instance = StructuredOutputsConfig()
-
-        engine_args = OmniEngineArgs(
-            model=model,
-            compilation_config=compilation_config_instance,
-            structured_outputs_config=structured_outputs_instance,
-            omni_kv_config=omni_kv_config,
-            hf_overrides=hf_overrides or {},
-            **filter_dataclass_kwargs(OmniEngineArgs, kwargs),
-        )
-
-        # Create the Engine (autoselects V0 vs V1)
-        self.llm_engine = LLMEngine.from_engine_args(engine_args=engine_args, usage_context=UsageContext.LLM_CLASS)
-        self.llm_engine.output_processor = MultimodalOutputProcessor(
-            tokenizer=self.llm_engine.tokenizer,
-            log_stats=self.llm_engine.log_stats,
-            engine_core_output_type=engine_args.engine_output_type,
-        )
-        self.llm_engine.input_processor = OmniInputProcessor(vllm_config=self.llm_engine.vllm_config)
-        self.engine_class = type(self.llm_engine)
-
-        self.request_counter = Counter()
-        self.default_sampling_params: dict[str, Any] | None = None
-
-        supported_tasks = self.llm_engine.get_supported_tasks()  # type: ignore
-
-        logger.info("Supported_tasks: %s", supported_tasks)
-
-        self.supported_tasks = supported_tasks
-
-        # Keep parity with vLLM's LLM initialization fields so inherited
-        # generate/chat preprocessing paths work as expected.
-        self.renderer = self.llm_engine.renderer
-        # Load the Input/Output processor plugin if any
-        io_processor_plugin = self.llm_engine.model_config.io_processor_plugin
-        self.io_processor = get_io_processor(self.llm_engine.vllm_config, io_processor_plugin)
-        self.model_config = self.llm_engine.model_config
-        self.input_processor = self.llm_engine.input_processor
-
-        # Parity with upstream LLM for pooling/classify entrypoints
-        chat_template = kwargs.get("chat_template", None)
-        from vllm.entrypoints.chat_utils import ChatTemplateConfig, load_chat_template
-        from vllm.entrypoints.pooling.io_processor_factories import init_pooling_io_processors
-
-        self.chat_template = load_chat_template(chat_template)
-        self.chat_template_config = ChatTemplateConfig(chat_template=self.chat_template)
-        self.init_pooling_io_processors = init_pooling_io_processors(
-            supported_tasks=supported_tasks,
-            model_config=self.model_config,
-            renderer=self.renderer,
-            chat_template_config=self.chat_template_config,
-        )
-
-    # ------------------------------------------------------------------
-    # Override upstream _preprocess_cmpl so that omni-specific fields
-    # (additional_information, prompt_embeds, …) survive the renderer's
-    # process_for_engine step which only copies standard vLLM keys.
-    # ------------------------------------------------------------------
-
-    def _preprocess_cmpl(
-        self,
-        prompts: Sequence[PromptType],
-        tokenization_kwargs: dict[str, Any] | None = None,
-    ) -> Sequence[ProcessorInputs]:
-        renderer = self.renderer
-        model_config = self.model_config
-
-        parsed_prompts = [parse_model_prompt(model_config, prompt) for prompt in prompts]
-        tok_params = renderer.default_cmpl_tok_params.with_kwargs(**(tokenization_kwargs or {}))
-        results = renderer.render_cmpl(parsed_prompts, tok_params)
-
-        reinject_omni_fields(results, parsed_prompts)
-        return results
-
-    def close(self) -> None:
-        """Close resources.
-
-        Note: Stage management is now handled by Omni class.
-        This method closes the LLM engine but not stages.
-        """
-        # Close the LLM engine if it exists
-        if hasattr(self, "llm_engine") and self.llm_engine is not None:
-            if hasattr(self.llm_engine, "shutdown"):
-                self.llm_engine.shutdown()
-
-    def __del__(self) -> None:  # best-effort
-        try:
-            self.close()
-        except Exception as e:
-            logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True)
-
-    def _run_engine(
-        self, output_type=None, *, use_tqdm: bool | Callable[..., tqdm] = True
-    ) -> list[RequestOutput | PoolingRequestOutput]:
-        # Initialize tqdm.
-        if use_tqdm:
-            num_requests = self.llm_engine.get_num_unfinished_requests()
-            tqdm_func = use_tqdm if callable(use_tqdm) else tqdm
-            pbar = tqdm_func(
-                total=num_requests,
-                desc="Processed prompts",
-                dynamic_ncols=True,
-                postfix=(f"est. speed input: {0:.2f} toks/s, output: {0:.2f} toks/s"),
-            )
-
-        # Run the engine.
-        outputs: list[RequestOutput | PoolingRequestOutput] = []
-        total_in_toks = 0
-        total_out_toks = 0
-        while self.llm_engine.has_unfinished_requests():
-            step_outputs = self.llm_engine.step()
-            for output in step_outputs:
-                if output.finished:
-                    outputs.append(output)
-                    if use_tqdm:
-                        if isinstance(output, RequestOutput):
-                            # Calculate tokens only for RequestOutput
-                            n = len(output.outputs)
-                            assert output.prompt_token_ids is not None
-                            total_in_toks += len(output.prompt_token_ids) * n
-                            in_spd = total_in_toks / pbar.format_dict["elapsed"]
-                            total_out_toks += sum(len(stp.token_ids) for stp in output.outputs)
-                            out_spd = total_out_toks / pbar.format_dict["elapsed"]
-                            pbar.postfix = f"est. speed input: {in_spd:.2f} toks/s, output: {out_spd:.2f} toks/s"
-                            pbar.update(n)
-                        else:
-                            pbar.update(1)
-                        if pbar.n == num_requests:
-                            pbar.refresh()
-
-        if use_tqdm:
-            pbar.close()
-        # Sort the outputs by the int part of request ID which is in format of 'int-uuid'.
-        # This is necessary because some requests may be finished earlier than
-        # its previous requests.
-        return sorted(outputs, key=lambda x: int(x.request_id.split("-")[0]))

vllm_omni/entrypoints/openai/api_server.py

@@ -111,10 +111,9 @@ from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpee
 from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage
 from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER
 from vllm_omni.entrypoints.openai.stores import VIDEO_STORE, VIDEO_TASKS
+from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request
 from vllm_omni.entrypoints.openai.video_api_utils import decode_input_reference
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams, OmniTextPrompt
-from vllm_omni.lora.request import LoRARequest
-from vllm_omni.lora.utils import stable_lora_int_id
 
 logger = init_logger(__name__)
 router = APIRouter()
@@ -163,6 +162,12 @@ def _remove_route_from_router(
 ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL = "endpoint-load-metrics-format"
 
 
+async def _get_vllm_config(engine_client: EngineClient) -> Any:
+    if hasattr(engine_client, "get_vllm_config"):
+        return await engine_client.get_vllm_config()
+    return getattr(engine_client, "vllm_config", None)
+
+
 def _remove_route_from_app(app, path: str, methods: set[str] | None = None):
     """Remove a route from the app by path and optionally by methods.
 
@@ -187,13 +192,31 @@ class _DiffusionServingModels:
     provide a lightweight fallback.
     """
 
+    class _NullModelConfig:
+        def __getattr__(self, name):
+            return None
+
+    class _Unsupported:
+        def __init__(self, name: str):
+            self.name = name
+
+        def __call__(self, *args, **kwargs):
+            raise NotImplementedError(f"{self.name} is not supported in diffusion mode")
+
+        def __getattr__(self, attr):
+            raise NotImplementedError(f"{self.name}.{attr} is not supported in diffusion mode")
+
     def __init__(self, base_model_paths: list[BaseModelPath]) -> None:
         self._base_model_paths = base_model_paths
+        self.model_config = self._NullModelConfig()
 
     @property
     def base_model_paths(self) -> list[BaseModelPath]:
         return self._base_model_paths
 
+    def __getattr__(self, name):
+        return self._Unsupported(name)
+
     async def show_available_models(self) -> ModelList:
         return ModelList(
             data=[
@@ -276,7 +299,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
             logger.warning("Profiler endpoints are enabled. This should ONLY be used for local development!")
             app.include_router(profiler_router)
 
-        vllm_config = await engine_client.get_vllm_config()
+        vllm_config = await _get_vllm_config(engine_client)
 
         # Check if pure diffusion mode (vllm_config will be None)
         is_pure_diffusion = vllm_config is None
@@ -334,7 +357,7 @@ async def build_async_omni(
     Args:
         args: Parsed command-line arguments containing model and configuration
         disable_frontend_multiprocessing: Optional flag to disable frontend
-            multiprocessing (deprecated in V1)
+            multiprocessing
         client_config: Optional client configuration dictionary
 
     Yields:
@@ -372,7 +395,7 @@ async def build_async_omni_from_stage_config(
     Args:
         args: Parsed command-line arguments containing model and stage configs
         disable_frontend_multiprocessing: Flag to disable frontend multiprocessing
-            (deprecated in V1)
+            for compatibility with existing CLI options
         client_config: Optional client configuration dictionary
 
     Yields:
@@ -383,16 +406,13 @@ async def build_async_omni_from_stage_config(
         otherwise from the model's default configuration.
     """
 
-    # V1 AsyncLLM.
     if disable_frontend_multiprocessing:
-        logger.warning("V1 is enabled, but got --disable-frontend-multiprocessing.")
+        logger.warning("Ignoring --disable-frontend-multiprocessing for AsyncOmni runtime.")
 
     async_omni: EngineClient | None = None
 
     try:
-        # Convert args Namespace to kwargs dict for AsyncOmni to use
         kwargs = vars(args).copy()
-        # Remove model as it will be passed separately
         kwargs.pop("model", None)
         async_omni = AsyncOmni(model=args.model, **kwargs)
 
@@ -423,14 +443,14 @@ async def omni_init_app_state(
         args: Parsed command-line arguments
     """
     # Get vllm_config from engine_client (following 0.14.0 pattern)
-    vllm_config = await engine_client.get_vllm_config()
+    vllm_config = await _get_vllm_config(engine_client)
 
     # Detect if it's pure Diffusion mode (single stage and is Diffusion)
     is_pure_diffusion = False
     if hasattr(engine_client, "stage_configs") and engine_client.stage_configs:
         stage_configs = engine_client.stage_configs
         if len(stage_configs) == 1:
-            stage_type = stage_configs[0].get("stage_type", "llm")
+            stage_type = get_stage_type(stage_configs[0])
             if stage_type == "diffusion":
                 is_pure_diffusion = True
                 logger.info("Detected pure diffusion mode (single diffusion stage)")
@@ -483,7 +503,7 @@ async def omni_init_app_state(
     # LLM or multi-stage mode: use standard initialization logic
     if vllm_config is None:
         # Try to get vllm_config from engine_client
-        vllm_config = await engine_client.get_vllm_config()
+        vllm_config = await _get_vllm_config(engine_client)
         if vllm_config is None:
             logger.warning("vllm_config is None, some features may not work correctly")
 
@@ -528,15 +548,13 @@ async def omni_init_app_state(
             # Try to initialize processors if vllm_config is available
             try:
                 from vllm.plugins.io_processors import get_io_processor
-
-                from vllm_omni.engine.input_processor import OmniInputProcessor
+                from vllm.v1.engine.input_processor import InputProcessor
 
                 tokenizer = await engine_client.get_tokenizer()
                 if tokenizer is not None:
                     # Initialize input_processor
-                    # OMNI: OmniInputProcessor creates tokenizer internally from vllm_config
                     if not hasattr(engine_client, "input_processor") or engine_client.input_processor is None:
-                        engine_client.input_processor = OmniInputProcessor(
+                        engine_client.input_processor = InputProcessor(
                             vllm_config=vllm_config,
                         )
                         logger.info("Initialized input_processor for AsyncOmni")
@@ -1167,7 +1185,7 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
         HTTPException: For validation errors, missing engine, or generation failures
     """
     # Get engine client (AsyncOmni) from app state
-    engine_client, model_name, stage_types = _get_engine_and_model(raw_request)
+    engine_client, model_name, stage_configs = _get_engine_and_model(raw_request)
 
     # Validate model field (warn if mismatch, don't error)
     if request.model is not None and request.model != model_name:
@@ -1219,7 +1237,7 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
         result = await _generate_with_async_omni(
             engine_client=engine_client,
             gen_params=gen_params,
-            stage_types=stage_types,
+            stage_configs=stage_configs,
             prompt=prompt,
             request_id=request_id,
         )
@@ -1294,7 +1312,7 @@ async def edit_images(
     OpenAI-compatible image edit endpoint.
     """
     # 1. get engine and model
-    engine_client, model_name, stage_types = _get_engine_and_model(raw_request)
+    engine_client, model_name, stage_configs = _get_engine_and_model(raw_request)
     if model is not None and model != model_name:
         logger.warning(
             f"Model mismatch: request specifies '{model}' but server is running '{model_name}'. Using server model."
@@ -1329,8 +1347,14 @@ async def edit_images(
         # 3.0 Init with system default values
         app_state_args = getattr(raw_request.app.state, "args", None)
         default_sample_param = getattr(app_state_args, "default_sampling_params", None)
-        # Currently only have one diffusion stage
-        diffusion_stage_id = [i for i, t in enumerate(stage_types) if t == "diffusion"][0]
+        # Currently only have one diffusion stage.
+        diffusion_stage_ids = [i for i, cfg in enumerate(stage_configs) if get_stage_type(cfg) == "diffusion"]
+        if not diffusion_stage_ids:
+            raise HTTPException(
+                status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
+                detail="No diffusion stage found in multi-stage pipeline.",
+            )
+        diffusion_stage_id = diffusion_stage_ids[0]
         apply_stage_default_sampling_params(
             default_sample_param,
             gen_params,
@@ -1377,7 +1401,7 @@ async def edit_images(
         result = await _generate_with_async_omni(
             engine_client=engine_client,
             gen_params=gen_params,
-            stage_types=stage_types,
+            stage_configs=stage_configs,
             prompt=prompt,
             request_id=request_id,
         )
@@ -1418,42 +1442,26 @@ async def edit_images(
 def _get_engine_and_model(raw_request: Request):
     # Get engine client (AsyncOmni) from app state
     engine_client: EngineClient | AsyncOmni | None = getattr(raw_request.app.state, "engine_client", None)
-    if engine_client is None or not hasattr(engine_client, "stage_list"):
+    if engine_client is None:
         raise HTTPException(
             status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
             detail="Multi-stage engine not initialized. Start server with a multi-stage omni model.",
         )
 
-    # Check if there's a diffusion stage
+    # Check if there's a diffusion stage.
+    # Prefer app state (compat layer populated at startup), then fall back to
+    # the engine client's stage configs for refactored AsyncOmni paths.
     stage_configs = getattr(raw_request.app.state, "stage_configs", None)
+    if not stage_configs:
+        stage_configs = getattr(engine_client, "stage_configs", None)
     if not stage_configs:
         raise HTTPException(
             status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
             detail="Stage configs not found. Start server with a multi-stage omni model.",
         )
 
-    # Check for diffusion stage and collect stage types
-    has_diffusion_stage = False
-    stage_types: list[str] = []
-    for stage in stage_configs:
-        # Handle both dict and OmegaConf objects
-        stage_type = None
-        if isinstance(stage, dict):
-            stage_type = stage.get("stage_type", "llm")
-        elif hasattr(stage, "get"):
-            stage_type = stage.get("stage_type", "llm")
-        elif hasattr(stage, "stage_type"):
-            stage_type = stage.stage_type
-        else:
-            # Fallback: try to access as dict-like
-            try:
-                stage_type = stage["stage_type"] if "stage_type" in stage else "llm"
-            except (TypeError, KeyError):
-                stage_type = "llm"
-
-        if stage_type == "diffusion":
-            has_diffusion_stage = True
-        stage_types.append(stage_type)
+    normalized_stage_configs = list(stage_configs)
+    has_diffusion_stage = any(get_stage_type(stage_cfg) == "diffusion" for stage_cfg in normalized_stage_configs)
 
     if not has_diffusion_stage:
         raise HTTPException(
@@ -1463,12 +1471,13 @@ def _get_engine_and_model(raw_request: Request):
 
     # Get server's loaded model name
     serving_models = getattr(raw_request.app.state, "openai_serving_models", None)
-    if serving_models and hasattr(serving_models, "base_model_paths") and serving_models.base_model_paths:
-        model_name = serving_models.base_model_paths[0].name
+    base_model_paths = getattr(serving_models, "base_model_paths", None) if serving_models else None
+    if base_model_paths:
+        model_name = base_model_paths[0].name
     else:
         model_name = "unknown"
 
-    return engine_client, model_name, stage_types
+    return engine_client, model_name, normalized_stage_configs
 
 
 def _get_lora_from_json_str(lora_body):
@@ -1486,81 +1495,63 @@ def _get_lora_from_json_str(lora_body):
 
 
 def _parse_lora_request(lora_body: dict[str, Any]):
-    if lora_body is not None:
-        if not isinstance(lora_body, dict):
-            raise HTTPException(
-                status_code=HTTPStatus.BAD_REQUEST.value,
-                detail="Invalid lora field: expected an object.",
-            )
-        lora_name = lora_body.get("name") or lora_body.get("lora_name") or lora_body.get("adapter")
-        lora_path = (
-            lora_body.get("local_path")
-            or lora_body.get("path")
-            or lora_body.get("lora_path")
-            or lora_body.get("lora_local_path")
-        )
-        lora_scale = lora_body.get("scale")
-        if lora_scale is None:
-            lora_scale = lora_body.get("lora_scale")
-        lora_int_id = lora_body.get("int_id")
-        if lora_int_id is None:
-            lora_int_id = lora_body.get("lora_int_id")
-        if lora_int_id is None and lora_path:
-            lora_int_id = stable_lora_int_id(str(lora_path))
-
-        if not lora_name or not lora_path:
-            raise HTTPException(
-                status_code=HTTPStatus.BAD_REQUEST.value,
-                detail="Invalid lora object: both name and path are required.",
-            )
-
-        return LoRARequest(str(lora_name), int(lora_int_id), str(lora_path)), lora_scale
-    return None, None
+    try:
+        return parse_lora_request(lora_body)
+    except ValueError as e:
+        raise HTTPException(
+            status_code=HTTPStatus.BAD_REQUEST.value,
+            detail=str(e),
+        ) from e
 
 
 async def _generate_with_async_omni(
     engine_client: AsyncOmni | Any,
     gen_params: Any,
-    stage_types: list[str],
+    stage_configs: list[Any],
     **kwargs,
 ):
     engine_client = cast(AsyncOmni, engine_client)
     result = None
-    stage_list = getattr(engine_client, "stage_list", None)
-    if isinstance(stage_list, list):
-        default_params_list: list[OmniSamplingParams] | None = getattr(
-            engine_client, "default_sampling_params_list", None
+    normalized_stage_configs = list(stage_configs)
+    if not normalized_stage_configs:
+        raise HTTPException(
+            status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
+            detail="Stage configs not found. Start server with a multi-stage omni model.",
         )
-        if not isinstance(default_params_list, list):
-            default_params_list = [
-                OmniDiffusionSamplingParams() if st == "diffusion" else SamplingParams() for st in stage_types
-            ]
-        else:
-            default_params_list = list(default_params_list)
-        if len(default_params_list) != len(stage_types):
-            default_params_list = (
-                default_params_list
-                + [OmniDiffusionSamplingParams() if st == "diffusion" else SamplingParams() for st in stage_types]
-            )[: len(stage_types)]
-
-        sampling_params_list: list[OmniSamplingParams] = []
-        for idx, stage_type in enumerate(stage_types):
-            if stage_type == "diffusion":
-                sampling_params_list.append(gen_params)
-            else:
-                base_params = default_params_list[idx]
-                sampling_params_list.append(base_params)
-
-        async for output in engine_client.generate(
-            sampling_params_list=sampling_params_list,
-            **kwargs,
-        ):
-            result = output
+    default_params_list: list[OmniSamplingParams] | None = getattr(
+        engine_client,
+        "default_sampling_params_list",
+        None,
+    )
+    if not isinstance(default_params_list, list):
+        default_params_list = []
     else:
-        result = await engine_client.generate(
-            sampling_params_list=[gen_params],
-            **kwargs,
-        )
+        default_params_list = list(default_params_list)
+
+    sampling_params_list: list[OmniSamplingParams] = []
+    for idx, stage_cfg in enumerate(normalized_stage_configs):
+        stage_type = get_stage_type(stage_cfg)
+        if stage_type == "diffusion":
+            sampling_params_list.append(gen_params)
+            continue
+
+        if idx < len(default_params_list):
+            default_stage_params = default_params_list[idx]
+        else:
+            default_stage_params = SamplingParams()
+
+        if hasattr(default_stage_params, "clone"):
+            try:
+                default_stage_params = default_stage_params.clone()
+            except Exception:
+                pass
+        sampling_params_list.append(default_stage_params)
+
+    async for output in engine_client.generate(
+        sampling_params_list=sampling_params_list,
+        **kwargs,
+    ):
+        result = output
 
     if result is None:
         raise HTTPException(

vllm_omni/entrypoints/openai/serving_chat.py

@@ -84,8 +84,8 @@ from vllm.utils.collection_utils import as_list
 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
 from vllm_omni.entrypoints.openai.protocol import OmniChatCompletionStreamResponse
 from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio
+from vllm_omni.entrypoints.openai.utils import parse_lora_request
 from vllm_omni.lora.request import LoRARequest
-from vllm_omni.lora.utils import stable_lora_int_id
 from vllm_omni.outputs import OmniRequestOutput
 
 if TYPE_CHECKING:
@@ -326,6 +326,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 tprompt: OmniTextPrompt = {"prompt": extracted_prompt}
                 if is_img2img:
                     tprompt["modalities"] = ["img2img"]
+                else:
+                    tprompt["modalities"] = ["image"]
                 if negative_prompt is not None:
                     tprompt["negative_prompt"] = negative_prompt
                 # GLM-Image's _call_hf_processor expects target_h/target_w in mm_processor_kwargs
@@ -626,10 +628,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         return final_sampling_params_list
 
     def _get_comprehension_stage_index(self) -> int:
-        for idx, stage in enumerate(self.engine_client.stage_list):
+        for idx, stage in enumerate(self.engine_client.stage_configs):
             if stage.is_comprehension:
                 return idx
-        raise ValueError("No comprehension stage (is_comprehension=True) found in stage_list")
+        raise ValueError("No comprehension stage (is_comprehension=True) found in stage configs")
 
     # OpenAI API standard sampling parameters that can be safely overridden.
     # These are the most commonly used parameters with compatible types
@@ -2104,27 +2106,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             # Parse per-request LoRA (works for both AsyncOmniDiffusion and AsyncOmni).
             if lora_body and isinstance(lora_body, dict):
                 try:
-                    lora_name = lora_body.get("name") or lora_body.get("lora_name") or lora_body.get("adapter")
-                    lora_path = (
-                        lora_body.get("local_path")
-                        or lora_body.get("path")
-                        or lora_body.get("lora_path")
-                        or lora_body.get("lora_local_path")
-                    )
-                    # using "or" directly here may be buggy if `scale=0`
-                    lora_scale = lora_body.get("scale")
-                    if lora_scale is None:
-                        lora_scale = lora_body.get("lora_scale")
-                    lora_int_id = lora_body.get("int_id")
-                    if lora_int_id is None:
-                        lora_int_id = lora_body.get("lora_int_id")
-                    if lora_int_id is None and lora_path:
-                        lora_int_id = stable_lora_int_id(str(lora_path))
-                    if lora_name and lora_path:
-                        lora_req = LoRARequest(str(lora_name), int(lora_int_id), str(lora_path))
+                    lora_req, lora_scale = parse_lora_request(lora_body)
+                    if lora_req is not None:
                         gen_params.lora_request = lora_req
                         if lora_scale is not None:
-                            gen_params.lora_scale = float(lora_scale)
+                            gen_params.lora_scale = lora_scale
                 except Exception as e:  # pragma: no cover - safeguard
                     logger.warning("Failed to parse LoRA request: %s", e)
 
@@ -2152,8 +2138,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
             # Generate image
             # Handle both AsyncOmniDiffusion (returns OmniRequestOutput) and AsyncOmni (returns AsyncGenerator)
-            if hasattr(self._diffusion_engine, "stage_list"):
-                # AsyncOmni: iterate through async generator to get final output
+            if isinstance(self._diffusion_engine, AsyncOmni):
                 diffusion_engine = cast(AsyncOmni, self._diffusion_engine)
                 result = None
                 async for output in diffusion_engine.generate(

vllm_omni/entrypoints/openai/serving_speech.py

@@ -197,12 +197,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         return None
 
     def _find_tts_stage(self):
-        """Find and return the TTS stage from the stage list, or None if not found."""
-        stage_list = getattr(self.engine_client, "stage_list", None)
-        if stage_list is None:
-            return None
-        for stage in stage_list:
-            if getattr(stage, "model_stage", None) in _TTS_MODEL_STAGES:
+        """Find and return the TTS stage config, or None if not found."""
+        for stage in self.engine_client.stage_configs:
+            if stage.engine_args.model_stage in _TTS_MODEL_STAGES:
                 return stage
         return None
 
@@ -496,13 +493,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
     def _is_tts_model(self) -> bool:
         """Check if the current model is a supported TTS model."""
-        stage_list = getattr(self.engine_client, "stage_list", None)
-        if stage_list:
-            for stage in stage_list:
-                model_stage = getattr(stage, "model_stage", None)
-                if model_stage in _TTS_MODEL_STAGES:
-                    return True
-        return False
+        return any(stage.engine_args.model_stage in _TTS_MODEL_STAGES for stage in self.engine_client.stage_configs)
 
     def _validate_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None:
         """Validate TTS request parameters. Returns error message or None."""