update

added omni mooncake connectors to distributed directory, integrate omni connector to omni-llm, rebased with main (twice).
ray + mooncake connector work on 2 nodes, but needes several updates\n1.start ray inside vllm-omni \n2.placement group optimize \n3.fix close
2025-12-01 10:49:39 +08:00 · 2025-12-01 10:45:26 +08:00
64 changed files with 2985 additions and 703 deletions
@@ -232,3 +232,4 @@ configs/development.yaml
 # Docker
 .dockerignore
 Dockerfile.dev
+discussion
@@ -1,3 +1,3 @@
 # Contributing to vLLM-Omni

-You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/)
+You may find information about contributing to vLLM-Omni on [Contributing](docs/contributing/)
@@ -1,7 +1,7 @@
 <p align="center">
  <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
-    <img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
+    <source media="(prefers-color-scheme: dark)" src="./docs/source/logos/vllm-omni-logo-text-dark.png">
+    <img alt="vllm-omni" src="./docs/source/logos/vllm-omni-logo-text-dark.png" width=55%>
  </picture>
 </p>
 <h3 align="center">
@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>

 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://tinyurl.com/vllm-omni-meeting"><b>Weekly Meeting</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> |
 </p>

 ---
@@ -30,7 +30,7 @@ Easy, fast, and cheap omni-modality model serving for everyone

 <p align="center">
  <picture>
-    <img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/omni-modality-model-architecture.png" width=55%>
+    <img alt="vllm-omni" src="./docs/source/architecture/omni-modality-model-architecture.png" width=55%>
  </picture>
 </p>

@@ -66,8 +66,10 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn
 We welcome and value any contributions and collaborations.
 Please check out [Contributing to vLLM-Omni](https://vllm-omni.readthedocs.io/en/latest/contributing/) for how to get involved.

-## Join the Community
-Feel free to ask questions, provide feedbacks and discuss with fellow users of vLLM-Omni in `#sig-omni` slack channel at [slack.vllm.ai](https://slack.vllm.ai) or vLLM user forum at [discuss.vllm.ai](https://discuss.vllm.ai).
+## Weekly Meeting
+
+- vLLM Omni Weekly Meeting: https://tinyurl.com/vllm-omni-meeting
+- Wednesday, 11:30 - 12:30 (UTC+8, [Convert to your timezone](https://dateful.com/convert/gmt8?t=15))

 ## License

@@ -7,11 +7,11 @@ nav:
  - Examples:
    - examples/README.md
    - Offline Inference:
-      - Offline Example of vLLM-Omni for Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
-      - Offline Example of vLLM-Omni for Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
+      - Offline Example of vLLM-Omni for Qwen2.5-omni: user_guide/examples/offline_inference/qwen2_5_omni.md
+      - Offline Example of vLLM-Omni for Qwen3-omni: user_guide/examples/offline_inference/qwen3_omni.md
      - Qwen-Image Offline Inference: user_guide/examples/offline_inference/qwen_image.md
    - Online Serving:
-      - Online serving Example of vLLM-Omni for Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
+      - Online serving Example of vLLM-Omni for Qwen2.5-omni: user_guide/examples/online_serving/qwen2_5_omni.md
      - Online serving Example of vLLM-Omni for Qwen3-omni: user_guide/examples/online_serving/qwen3_omni.md
  - General:
    - usage/*
@@ -33,6 +33,7 @@ nav:
    - design/architecture_overview.md
    - design/vllm_omni_design.md
    - design/mrs_design.md
+    - design/api_design_doc.md
  - Docs Guide: contributing/DOCS_GUIDE.md
 - API Reference:
  - api/README.md
@@ -8,8 +8,8 @@ hide:

 <p align="center">
  <picture>
-    <source media="(prefers-color-scheme: dark)" src="./source/logos/vllm-omni-logo.png">
-    <img alt="vllm-omni" src="./source/logos/vllm-omni-logo.png" width=55%>
+    <source media="(prefers-color-scheme: dark)" src="./source/logos/vllm-omni-logo-text-dark.png">
+    <img alt="vllm-omni" src="./source/logos/vllm-omni-logo-text-dark.png" width=55%>
  </picture>
 </p>
 <h3 align="center">
@@ -34,8 +34,7 @@ Easy, fast, and cheap omni-modality model serving for everyone

 <p align="center">
  <picture>
-    <source media="(prefers-color-scheme: dark)" src="./source/architecture/omni-modality-model-architecture.png">
-    <img alt="vllm-omni-arch" src="./source/architecture/omni-modality-model-architecture.png" width=55%>
+    <img alt="vllm-omni" src="./source/architecture/omni-modality model architecture.png" width=55%>
  </picture>
 </p>

@@ -59,4 +58,4 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
 - Multi-modality generation models (e.g. Qwen-Image)

 For more information, checkout the following:
- [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
+- [Placeholder for vllm-omni blogpost]()
@@ -102,10 +102,6 @@ Configuration classes.

 Worker classes and model runners for distributed inference.

- [vllm_omni.diffusion.worker.gpu_worker.GPUWorker][]
- [vllm_omni.diffusion.worker.gpu_worker.WorkerProc][]
- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorker][]
- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorkerProc][]
 - [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][]
 - [vllm_omni.worker.gpu_ar_worker.GPUARWorker][]
 - [vllm_omni.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][]
@@ -2,9 +2,6 @@

 In vLLM-Omni, the target model is separated into multiple stages, which are processed by different LLMEngines, DiffusionEngines or other types of engines. Depending on different types of stages, such as Autoregressive (AR) stage or Diffusion transformer (DiT) stage, each can choose corresponding schedulers, model workers to load with the Engines in a plug-in fashion.

-!!! note
-    Default stage config YAMLs (for example, `vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml` and `vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml`) are bundled and loaded automatically when `stage_configs_path` is not provided. They have been verified to work on 1xH100 for Qwen2.5-Omni and 2xH100 for Qwen3-Omni.
-
 Therefore, as a core part of vLLM-Omni, the stage configs for a model have several main functions:

 - Claim partition of stages and their corresponding class implementation in `model_executor/models`.
@@ -2,9 +2,6 @@

 Thank you for your interest in contributing to vLLM-Omni! This document provides guidelines and instructions for contributing.

-!!! note
-    We host weekly developer-facing online meetings to discuss milestones and updates **every Tuesday at 19:30 PDT**. Meeting link as well as the past meeting notes can be found [here](https://tinyurl.com/vllm-omni-meeting).
-
 ## Getting Started

 vLLM-Omni uses `uv` as the environment manager, to create and manage Python environments. Please follow the documentation to install `uv`. After installing `uv`, you can create a new Python environment using the following commands:
@@ -7,3 +7,7 @@ This section contains design documents and architecture specifications for vLLM-
 - [Architecture Overview](architecture_overview.md)
 - [vLLM-Omni Software Design Document](vllm_omni_design.md)
 - [Multi-Request Streaming](mrs_design.md)
+
+## API Design Documentation
+
+- [vLLM-Omni API Design Documentation](api_design_doc.md)
@@ -19,7 +19,6 @@ export DEVICE1=/dev/davinci1
 export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2
 docker run --rm \
    --name vllm-omni-npu \
-    --shm-size=1g \
    --device $DEVICE0 \
    --device $DEVICE1 \
    --device /dev/davinci_manager \
@@ -192,8 +192,7 @@ class Example:
        code_fence = "``````"

        if self.is_code:
-            main_file_rel = self.main_file.relative_to(ROOT_DIR)
-            content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{main_file_rel}"\n{code_fence}\n'
+            content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{self.main_file}"\n{code_fence}\n'
        else:
            with open(self.main_file) as f:
                # Skip the title from md snippets as it's been included above
@@ -209,7 +208,7 @@ class Example:
            content += f'??? abstract "{file.relative_to(self.path)}"\n'
            if file.suffix != ".md":
                content += f"    {code_fence}{file.suffix[1:]}\n"
-            content += f'    --8<-- "{file.relative_to(ROOT_DIR)}"\n'
+            content += f'    --8<-- "{file}"\n'
            if file.suffix != ".md":
                content += f"    {code_fence}\n"

@@ -1,4 +1,4 @@
-# Offline Example of vLLM-Omni for Qwen2.5-Omni
+# Offline Example of vLLM-Omni for Qwen2.5-omni

 Source <https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/qwen2_5_omni>.

@@ -7,7 +7,7 @@ Source <https://github.com/vllm-project/vllm/tree/main/examples/offline_inferenc

 Please refer to [README.md](https://github.com/vllm-project/vllm/tree/main/README.md)

-## Run examples (Qwen2.5-Omni)
+## Run examples (Qwen2.5-omni)
 ### Multiple Prompts
 Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
 ```bash
@@ -48,17 +48,17 @@ sudo apt install ffmpeg

 ??? abstract "end2end.py"
    ``````py
-    --8<-- "examples/offline_inference/qwen2_5_omni/end2end.py"
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen2_5_omni/end2end.py"
    ``````
 ??? abstract "extract_prompts.py"
    ``````py
-    --8<-- "examples/offline_inference/qwen2_5_omni/extract_prompts.py"
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen2_5_omni/extract_prompts.py"
    ``````
 ??? abstract "run_multiple_prompts.sh"
    ``````sh
-    --8<-- "examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh"
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen2_5_omni/run_multiple_prompts.sh"
    ``````
 ??? abstract "run_single_prompt.sh"
    ``````sh
-    --8<-- "examples/offline_inference/qwen2_5_omni/run_single_prompt.sh"
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen2_5_omni/run_single_prompt.sh"
    ``````
@@ -1,4 +1,4 @@
-# Offline Example of vLLM-Omni for Qwen3-Omni
+# Offline Example of vLLM-Omni for Qwen3-omni

 Source <https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/qwen3_omni>.

@@ -7,9 +7,9 @@ Source <https://github.com/vllm-project/vllm/tree/main/examples/offline_inferenc

 Please refer to [README.md](https://github.com/vllm-project/vllm/tree/main/README.md)

-## Run examples (Qwen3-Omni)
+## Run examples (Qwen3-omni)
 ### Multiple Prompts
-Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-Omni README.md](https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/qwen2_5_omni/README.md)
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-omni README.md](https://github.com/vllm-project/vllm/tree/main/examples/offline_inference/qwen2_5_omni/README.md)
 Get into the example folder
 ```bash
 cd examples/offline_inference/qwen3_omni
@@ -44,17 +44,21 @@ sudo apt install ffmpeg

 ??? abstract "end2end.py"
    ``````py
-    --8<-- "examples/offline_inference/qwen3_omni/end2end.py"
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/end2end.py"
+    ``````
+??? abstract "qwen3_omni_moe_tp.yaml"
+    ``````yaml
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/qwen3_omni_moe_tp.yaml"
    ``````
 ??? abstract "run_multiple_prompts.sh"
    ``````sh
-    --8<-- "examples/offline_inference/qwen3_omni/run_multiple_prompts.sh"
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/run_multiple_prompts.sh"
    ``````
 ??? abstract "run_single_prompt.sh"
    ``````sh
-    --8<-- "examples/offline_inference/qwen3_omni/run_single_prompt.sh"
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/run_single_prompt.sh"
    ``````
 ??? abstract "run_single_prompt_tp.sh"
    ``````sh
-    --8<-- "examples/offline_inference/qwen3_omni/run_single_prompt_tp.sh"
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen3_omni/run_single_prompt_tp.sh"
    ``````
@@ -49,9 +49,9 @@ Then open `http://localhost:7862/` on your local browser to interact with the we

 ??? abstract "gradio_demo.py"
    ``````py
-    --8<-- "examples/offline_inference/qwen_image/gradio_demo.py"
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen_image/gradio_demo.py"
    ``````
 ??? abstract "text_to_image.py"
    ``````py
-    --8<-- "examples/offline_inference/qwen_image/text_to_image.py"
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/offline_inference/qwen_image/text_to_image.py"
    ``````
@@ -1,13 +1,13 @@
-# Online serving Example of vLLM-Omni for Qwen2.5-Omni
+# Online serving Example of vLLM-Omni for Qwen2.5-omni

 Source <https://github.com/vllm-project/vllm/tree/main/examples/online_serving/qwen2_5_omni>.


 ## 🛠️ Installation

-Please refer to [README.md](https://github.com/vllm-project/vllm/tree/main/README.md)
+Please refer to [README.md](https://github.com/vllm-project/vllm/tree/main/examples/README.md)

-## Run examples (Qwen2.5-Omni)
+## Run examples (Qwen2.5-omni)

 Launch the server
 ```bash
@@ -21,7 +21,7 @@ vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to

 Get into the example folder
 ```bash
-cd examples/online_serving/qwen2_5_omni
+cd examples/online_serving
 ```

 Send request via python
@@ -79,13 +79,13 @@ python gradio_demo.py \

 ??? abstract "gradio_demo.py"
    ``````py
-    --8<-- "examples/online_serving/qwen2_5_omni/gradio_demo.py"
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen2_5_omni/gradio_demo.py"
    ``````
 ??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
    ``````py
-    --8<-- "examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py"
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py"
    ``````
 ??? abstract "run_curl_multimodal_generation.sh"
    ``````sh
-    --8<-- "examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh"
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh"
    ``````
@@ -5,7 +5,7 @@ Source <https://github.com/vllm-project/vllm/tree/main/examples/online_serving/q

 ## 🛠️ Installation

-Please refer to [README.md](https://github.com/vllm-project/vllm/tree/main/README.md)
+Please refer to [README.md](https://github.com/vllm-project/vllm/tree/main/examples/README.md)

 ## Run examples (Qwen3-Omni)

@@ -21,7 +21,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-p

 Get into the example folder
 ```bash
-cd examples/online_serving/qwen3_omni
+cd examples/online_serving
 ```

 Send request via python
@@ -83,13 +83,13 @@ python gradio_demo.py \

 ??? abstract "gradio_demo.py"
    ``````py
-    --8<-- "examples/online_serving/qwen3_omni/gradio_demo.py"
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen3_omni/gradio_demo.py"
    ``````
 ??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
    ``````py
-    --8<-- "examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py"
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py"
    ``````
 ??? abstract "run_curl_multimodal_generation.sh"
    ``````sh
-    --8<-- "examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh"
+    --8<-- "/mnt/vllm_open_release/vllm-omni-cursor/vllm-omni/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh"
    ``````
@@ -283,6 +283,19 @@ def parse_args():
        default=None,
        help="Path to a .txt file with one prompt per line (preferred).",
    )
+    parser.add_argument(
+        "--worker-backend",
+        type=str,
+        default="process",
+        choices=["process","ray"],
+        help="backend"
+    )
+    parser.add_argument(
+        "--ray-address",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred).",
+

    return parser.parse_args()

@@ -0,0 +1,94 @@
+# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
+# Stage 0: Thinker (multimodal understanding + text generation)
+# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes)
+# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
+
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0,1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      gpu_memory_utilization: 0.8
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent  # Output hidden states for talker
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      hf_config_name: thinker_config
+      tensor_parallel_size: 2
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    runtime:
+       devices: "1"
+       max_batch_size: 1
+    engine_args:
+       model_stage: talker
+       model_arch: Qwen3OmniMoeForConditionalGeneration
+       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+       gpu_memory_utilization: 0.5
+       enforce_eager: true
+       trust_remote_code: true
+       engine_output_type: latent  # Output codec codes for code2wav
+      #  tensor_parallel_size: 2
+       enable_prefix_caching: false
+       distributed_executor_backend: "mp"
+       hf_config_name: talker_config
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+    # final_output: true
+    # final_output_type: text
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: False
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 2
+    runtime:
+      devices: "1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.generation_scheduler.GenerationScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: audio  # Final output: audio waveform
+      gpu_memory_utilization: 0.1
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 1000000
+      hf_config_name: thinker_config
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
+    final_output: true
+    final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
@@ -1,5 +1,6 @@
 python end2end.py --output-wav output_audio \
                  --query-type use_audio \
-                  --init-sleep-seconds 90
+                  --init-sleep-seconds 90 \
+                  --stage-configs-path qwen3_omni_moe_tp.yaml

 # init-sleep-seconds works to avoid two vLLM stages initialized at the same time within a card.
@@ -5,7 +5,7 @@ import gradio as gr
 import torch

 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.utils.platform_utils import detect_device_type

 ASPECT_RATIOS: dict[str, tuple[int, int]] = {
    "1:1": (1328, 1328),
@@ -59,14 +59,7 @@ def parse_args() -> argparse.Namespace:

@lru_cache(maxsize=1)
 def get_omni(model_name: str) -> Omni:
-    # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
-    return Omni(
-        model=model_name,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
-    )
+    return Omni(model=model_name)


 def build_demo(args: argparse.Namespace) -> gr.Blocks:
@@ -7,7 +7,7 @@ from pathlib import Path
 import torch

 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.utils.platform_utils import detect_device_type


 def parse_args() -> argparse.Namespace:
@@ -49,15 +49,7 @@ def main():
    device = detect_device_type()
    generator = torch.Generator(device=device).manual_seed(args.seed)

-    # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
-
-    omni = Omni(
-        model=args.model,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
-    )
+    omni = Omni(model=args.model)
    images = omni.generate(
        args.prompt,
        height=args.height,
@@ -0,0 +1,34 @@
+# Offline Example of vLLM-omni for Qwen2.5-omni
+
+## 🛠️ Installation
+
+Please refer to [README.md](../../../README.md)
+
+## Run examples (Qwen2.5-omni)
+### Multiple Prompts
+Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). To get the prompt, you can:
+```bash
+tar -xf <Your Download Path>/seedtts_testset.tar
+cp seedtts_testset/en/meta.lst examples/offline_inference/qwen2_5_omni/meta.lst
+python3 examples/offline_inference/qwen2_5_omni/extract_prompts.py \
+  --input examples/offline_inference/qwen2_5_omni/meta.lst \
+  --output examples/offline_inference/qwen2_5_omni/top100.txt \
+  --topk 100
+```
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen2_5_omni
+```
+Then run the command below.
+```bash
+bash run_multiple_prompts.sh
+```
+### Single Prompts
+Get into the example folder
+```bash
+cd examples/offline_inference/qwen2_5_omni
+```
+Then run the command below.
+```bash
+bash run_single_prompt.sh
+```
@@ -0,0 +1,264 @@
+import argparse
+import os
+import os as _os_env_toggle
+import random
+
+import numpy as np
+import soundfile as sf
+import torch
+from utils import make_omni_prompt
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+_os_env_toggle.environ["VLLM_USE_V1"] = "1"
+
+SEED = 42
+# Set all random seeds
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+torch.cuda.manual_seed(SEED)
+torch.cuda.manual_seed_all(SEED)
+
+# Make PyTorch deterministic
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+
+# Set environment variables for deterministic behavior
+os.environ["PYTHONHASHSEED"] = str(SEED)
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model",
+        required=True,
+        help="Path to merged model directory (will be created if downloading).",
+    )
+    parser.add_argument("--thinker-model", type=str, default=None)
+    parser.add_argument("--talker-model", type=str, default=None)
+    parser.add_argument("--code2wav-model", type=str, default=None)
+    parser.add_argument(
+        "--hf-hub-id",
+        default="Qwen/Qwen2.5-Omni-7B",
+        help="Hugging Face repo id to download if needed.",
+    )
+    parser.add_argument("--hf-revision", default=None, help="Optional HF revision (branch/tag/commit).")
+    parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
+    parser.add_argument("--voice-type", default="default", help="Voice type, e.g., m02, f030, default.")
+    parser.add_argument(
+        "--code2wav-dir",
+        default=None,
+        help="Path to code2wav folder (contains spk_dict.pt).",
+    )
+    parser.add_argument("--dit-ckpt", default=None, help="Path to DiT checkpoint file (e.g., dit.pt).")
+    parser.add_argument("--bigvgan-ckpt", default=None, help="Path to BigVGAN checkpoint file.")
+    parser.add_argument("--dtype", default="bfloat16", choices=["float16", "bfloat16", "float32"])
+    parser.add_argument("--max-model-len", type=int, default=32768)
+    parser.add_argument(
+        "--init-sleep-seconds",
+        type=int,
+        default=20,
+        help="Sleep seconds after starting each stage process to allow initialization (default: 20)",
+    )
+
+    parser.add_argument("--thinker-only", action="store_true")
+    parser.add_argument("--text-only", action="store_true")
+    parser.add_argument("--do-wave", action="store_true")
+    parser.add_argument(
+        "--prompt_type",
+        choices=[
+            "text",
+            "audio",
+            "audio-long",
+            "audio-long-chunks",
+            "audio-long-expand-chunks",
+            "image",
+            "video",
+            "video-frames",
+            "audio-in-video",
+            "audio-in-video-v2",
+            "audio-multi-round",
+            "badcase-vl",
+            "badcase-text",
+            "badcase-image-early-stop",
+            "badcase-two-audios",
+            "badcase-two-videos",
+            "badcase-multi-round",
+            "badcase-voice-type",
+            "badcase-voice-type-v2",
+            "badcase-audio-tower-1",
+            "badcase-audio-only",
+        ],
+        default="text",
+    )
+    parser.add_argument("--use-torchvision", action="store_true")
+    parser.add_argument("--tokenize", action="store_true")
+    parser.add_argument(
+        "--output-wav",
+        default="output.wav",
+        help="[Deprecated] Output wav directory (use --output-dir).",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default="outputs",
+        help="Output directory to save text and wav files together.",
+    )
+    parser.add_argument(
+        "--thinker-hidden-states-dir",
+        default="thinker_hidden_states",
+        help="Path to thinker hidden states directory.",
+    )
+    parser.add_argument(
+        "--batch-timeout",
+        type=int,
+        default=5,
+        help="Timeout for batching in seconds (default: 5)",
+    )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=300,
+        help="Timeout for initializing stages in seconds (default: 300)",
+    )
+    parser.add_argument(
+        "--shm-threshold-bytes",
+        type=int,
+        default=65536,
+        help="Threshold for using shared memory in bytes (default: 65536)",
+    )
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        default=False,
+        help="Enable writing detailed statistics (default: disabled)",
+    )
+    parser.add_argument(
+        "--txt-prompts",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred).",
+    )
+    parser.add_argument(
+        "--worker-backend",
+        type=str,
+        default="process",
+        choices=["process","ray"],
+        help="backend"
+    )
+    parser.add_argument(
+        "--ray-address",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred)."
+        )
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    model_name = args.model
+    try:
+        # Preferred: load from txt file (one prompt per line)
+        if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
+            with open(args.txt_prompts, encoding="utf-8") as f:
+                lines = [ln.strip() for ln in f.readlines()]
+            args.prompts = [ln for ln in lines if ln != ""]
+            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
+    except Exception as e:
+        print(f"[Error] Failed to load prompts: {e}")
+        raise
+
+    if args.prompts is None:
+        raise ValueError("No prompts provided. Use --prompts ... or --txt-prompts <file.txt> (with --prompt_type text)")
+    omni_llm = OmniLLM(
+        model=model_name,
+        log_stats=args.enable_stats,
+        log_file=("omni_llm_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,
+        shm_threshold_bytes=args.shm_threshold_bytes,
+        worker_backend=args.worker_backend,
+        ray_address=args.ray_address
+    )
+    thinker_sampling_params = SamplingParams(
+        temperature=0.0,  # Deterministic - no randomness
+        top_p=1.0,  # Disable nucleus sampling
+        top_k=-1,  # Disable top-k sampling
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.1,
+    )
+    talker_sampling_params = SamplingParams(
+        temperature=0.9,
+        top_p=0.8,
+        top_k=40,
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.05,
+        stop_token_ids=[8294],
+    )
+    code2wav_sampling_params = SamplingParams(
+        temperature=0.0,  # Deterministic - no randomness
+        top_p=1.0,  # Disable nucleus sampling
+        top_k=-1,  # Disable top-k sampling
+        max_tokens=2048,
+        seed=SEED,  # Fixed seed for sampling
+        detokenize=True,
+        repetition_penalty=1.1,
+    )
+
+    sampling_params_list = [
+        thinker_sampling_params,
+        talker_sampling_params,
+        code2wav_sampling_params,
+    ]
+    import time
+    for i in range(1):
+        t1 = time.time()
+        prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
+        print(f"prompt:{prompt}")
+        omni_outputs = omni_llm.generate(prompt, sampling_params_list)
+        t2 = time.time()
+        print(f"==========> time:{t2-t1}")
+
+    # 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)
+    for stage_outputs in omni_outputs:
+        if stage_outputs.final_output_type == "text":
+            for output in stage_outputs.request_output:
+                request_id = int(output.request_id)
+                text_output = output.outputs[0].text
+                # Save aligned text file per request
+                prompt_text = args.prompts[request_id]
+                out_txt = os.path.join(output_dir, f"{request_id:05d}.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 = int(output.request_id)
+                audio_tensor = output.multimodal_output["audio"]
+                output_wav = os.path.join(output_dir, f"output_{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}")
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+import argparse
+from typing import Optional
+
+
+def extract_prompt(line: str) -> Optional[str]:
+    # 提取第一个 '|' 与第二个 '|' 之间的内容
+    i = line.find("|")
+    if i == -1:
+        return None
+    j = line.find("|", i + 1)
+    if j == -1:
+        return None
+    return line[i + 1 : j].strip()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input", "-i", required=True, help="Input .lst file path")
+    parser.add_argument("--output", "-o", required=True, help="Output file path")
+    parser.add_argument(
+        "--topk",
+        "-k",
+        type=int,
+        default=100,
+        help="Extract the top K prompts (default: 100)",
+    )
+    args = parser.parse_args()
+
+    prompts = []
+    with open(args.input, encoding="utf-8", errors="ignore") as f:
+        for line in f:
+            if len(prompts) >= args.topk:
+                break
+            p = extract_prompt(line.rstrip("\n"))
+            if p:
+                prompts.append(p)
+
+    with open(args.output, "w", encoding="utf-8") as f:
+        for p in prompts:
+            f.write(p + "\n")
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,367 @@
+from __future__ import annotations
+
+import base64
+import logging
+import math
+import os
+import time
+import warnings
+from functools import lru_cache
+from io import BytesIO
+
+import requests
+import torch
+import torchvision
+from packaging import version
+from PIL import Image
+from torchvision import io, transforms
+from torchvision.transforms import InterpolationMode
+
+logger = logging.getLogger(__name__)
+
+IMAGE_FACTOR = 28
+MIN_PIXELS = 4 * 28 * 28
+MAX_PIXELS = 16384 * 28 * 28
+MAX_RATIO = 200
+
+VIDEO_MIN_PIXELS = 128 * 28 * 28
+VIDEO_MAX_PIXELS = 768 * 28 * 28
+VIDEO_TOTAL_PIXELS = 24576 * 28 * 28
+FRAME_FACTOR = 2
+FPS = 2.0
+FPS_MIN_FRAMES = 4
+FPS_MAX_FRAMES = 768
+
+temporal_patch_size = 2
+spatial_patch_size = 14
+spatial_merge_size = 2
+
+
+def round_by_factor(number: int, factor: int) -> int:
+    """Returns the closest integer to 'number' that is divisible by 'factor'."""
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: int, factor: int) -> int:
+    """Returns the smallest integer greater than or equal to 'number' that is
+    divisible by 'factor'."""
+    return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: int, factor: int) -> int:
+    """Returns the largest integer less than or equal to 'number' that is
+    divisible by 'factor'."""
+    return math.floor(number / factor) * factor
+
+
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int = IMAGE_FACTOR,
+    min_pixels: int = MIN_PIXELS,
+    max_pixels: int = MAX_PIXELS,
+) -> tuple[int, int]:
+    """
+    Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+    3. The aspect ratio of the image is maintained as closely as possible.
+    """
+    if max(height, width) / min(height, width) > MAX_RATIO:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(height / beta, factor)
+        w_bar = floor_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+    return h_bar, w_bar
+
+
+def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image:
+    if "image" in ele:
+        image = ele["image"]
+    else:
+        image = ele["image_url"]
+    image_obj = None
+    if isinstance(image, Image.Image):
+        image_obj = image
+    elif image.startswith("http://") or image.startswith("https://"):
+        image_obj = Image.open(requests.get(image, stream=True).raw)
+    elif image.startswith("file://"):
+        image_obj = Image.open(image[7:])
+    elif image.startswith("data:image"):
+        if "base64," in image:
+            _, base64_data = image.split("base64,", 1)
+            data = base64.b64decode(base64_data)
+            image_obj = Image.open(BytesIO(data))
+    else:
+        image_obj = Image.open(image)
+    if image_obj is None:
+        raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
+    image = image_obj.convert("RGB")
+    # resize
+    if "resized_height" in ele and "resized_width" in ele:
+        resized_height, resized_width = smart_resize(
+            ele["resized_height"],
+            ele["resized_width"],
+            factor=size_factor,
+        )
+    else:
+        width, height = image.size
+        min_pixels = ele.get("min_pixels", MIN_PIXELS)
+        max_pixels = ele.get("max_pixels", MAX_PIXELS)
+        resized_height, resized_width = smart_resize(
+            height,
+            width,
+            factor=size_factor,
+            min_pixels=min_pixels,
+            max_pixels=max_pixels,
+        )
+    image = image.resize((resized_width, resized_height))
+
+    return image
+
+
+def smart_nframes(
+    ele: dict,
+    total_frames: int,
+    video_fps: int | float,
+) -> int:
+    """calculate the number of frames for video used for model inputs.
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+            support either `fps` or `nframes`:
+                - nframes: the number of frames to extract for model inputs.
+                - fps: the fps to extract frames for model inputs.
+                    - min_frames: the minimum number of frames of the video,
+                      only used when fps is provided.
+                    - max_frames: the maximum number of frames of the video,
+                      only used when fps is provided.
+        total_frames (int): the original total number of frames of the video.
+        video_fps (int | float): the original fps of the video.
+
+    Raises:
+        ValueError: nframes should in interval [FRAME_FACTOR, total_frames].
+
+    Returns:
+        int: the number of frames for video used for model inputs.
+    """
+    assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`"
+    if "nframes" in ele:
+        nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
+    else:
+        fps = ele.get("fps", FPS)
+        min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR)
+        max_frames = floor_by_factor(
+            ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)),
+            FRAME_FACTOR,
+        )
+        nframes = total_frames / video_fps * fps
+        nframes = min(max(nframes, min_frames), max_frames)
+        nframes = round_by_factor(nframes, FRAME_FACTOR)
+    if not (FRAME_FACTOR <= nframes and nframes <= total_frames):
+        raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.")
+    return nframes
+
+
+def _read_video_torchvision(
+    ele: dict,
+) -> torch.Tensor:
+    """read video using torchvision.io.read_video
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+        support keys:
+            - video: the path of video. support "file://", "http://",
+              "https://" and local path.
+            - video_start: the start time of video.
+            - video_end: the end time of video.
+    Returns:
+        torch.Tensor: the video tensor with shape (T, C, H, W).
+    """
+    video_path = ele["video"]
+    if version.parse(torchvision.__version__) < version.parse("0.19.0"):
+        if "http://" in video_path or "https://" in video_path:
+            warnings.warn("torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.")
+        if "file://" in video_path:
+            video_path = video_path[7:]
+    st = time.time()
+    video, audio, info = io.read_video(
+        video_path,
+        start_pts=ele.get("video_start", 0.0),
+        end_pts=ele.get("video_end", None),
+        pts_unit="sec",
+        output_format="TCHW",
+    )
+    total_frames, video_fps = video.size(0), info["video_fps"]
+    total_duration = round(total_frames / video_fps, 3)
+    logger.info(
+        f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, duration={total_duration}s, time={time.time() - st:.3f}s"
+    )
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(0, total_frames - 1, nframes).round().long()
+    video = video[idx]
+    return video, total_duration, nframes
+
+
+def is_decord_available() -> bool:
+    import importlib.util
+
+    return importlib.util.find_spec("decord") is not None
+
+
+def _read_video_decord(
+    ele: dict,
+) -> torch.Tensor:
+    """read video using decord.VideoReader
+
+    Args:
+        ele (dict): a dict contains the configuration of video.
+        support keys:
+            - video: the path of video. support "file://", "http://",
+              "https://" and local path.
+            - video_start: the start time of video.
+            - video_end: the end time of video.
+    Returns:
+        torch.Tensor: the video tensor with shape (T, C, H, W).
+    """
+    import decord
+
+    video_path = ele["video"]
+    st = time.time()
+    vr = decord.VideoReader(video_path)
+    # TODO: support start_pts and end_pts
+    if "video_start" in ele or "video_end" in ele:
+        raise NotImplementedError("not support start_pts and end_pts in decord for now.")
+    total_frames, video_fps = len(vr), vr.get_avg_fps()
+    total_duration = round(total_frames / video_fps, 3)
+    logger.info(f"decord:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
+    video = vr.get_batch(idx).asnumpy()
+    video = torch.tensor(video).permute(0, 3, 1, 2)  # Convert to TCHW format
+    return video, total_duration, nframes
+
+
+VIDEO_READER_BACKENDS = {
+    "decord": _read_video_decord,
+    "torchvision": _read_video_torchvision,
+}
+
+FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None)
+
+
+@lru_cache(maxsize=1)
+def get_video_reader_backend() -> str:
+    if FORCE_QWENVL_VIDEO_READER is not None:
+        video_reader_backend = FORCE_QWENVL_VIDEO_READER
+    elif is_decord_available():
+        video_reader_backend = "decord"
+    else:
+        video_reader_backend = "torchvision"
+    # print(f"qwen-vl-utils using {video_reader_backend} to read video.",
+    #       file=sys.stderr)
+    return video_reader_backend
+
+
+def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR) -> torch.Tensor | list[Image.Image]:
+    if isinstance(ele["video"], str):
+        video_reader_backend = get_video_reader_backend()
+        video, total_dur, nframes = VIDEO_READER_BACKENDS[video_reader_backend](ele)
+        frame_timestamps = total_dur * torch.arange(1, nframes + 1) / nframes
+        grid_timestamps = frame_timestamps[::FRAME_FACTOR]
+        second_per_grid = grid_timestamps[1] - grid_timestamps[0]
+        nframes, _, height, width = video.shape
+        min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS)
+        total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS)
+        max_pixels = max(
+            min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR),
+            int(min_pixels * 1.05),
+        )
+        max_pixels = ele.get("max_pixels", max_pixels)
+        # min_pixels = (factor ** 2) * 52
+        # max_pixels = (factor ** 2) * min(768, (16384 / nframes * temporal_patch_size))
+        if "resized_height" in ele and "resized_width" in ele:
+            resized_height, resized_width = smart_resize(
+                ele["resized_height"],
+                ele["resized_width"],
+                factor=image_factor,
+            )
+        else:
+            resized_height, resized_width = smart_resize(
+                height,
+                width,
+                factor=image_factor,
+                min_pixels=min_pixels,
+                max_pixels=max_pixels,
+            )
+        video = transforms.functional.resize(
+            video,
+            [resized_height, resized_width],
+            interpolation=InterpolationMode.BICUBIC,
+            antialias=True,
+        ).float()
+        return video, total_dur, nframes, second_per_grid
+    else:
+        assert isinstance(ele["video"], (list, tuple))
+        process_info = ele.copy()
+        process_info.pop("type", None)
+        process_info.pop("video", None)
+        images = [
+            fetch_image({"image": video_element, **process_info}, size_factor=image_factor)
+            for video_element in ele["video"]
+        ]
+        nframes = ceil_by_factor(len(images), FRAME_FACTOR)
+        if len(images) < nframes:
+            images.extend([images[-1]] * (nframes - len(images)))
+        return images, None, None, None
+
+
+def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]:
+    vision_infos = []
+    if isinstance(conversations[0], dict):
+        conversations = [conversations]
+    for conversation in conversations:
+        for message in conversation:
+            if isinstance(message["content"], list):
+                for ele in message["content"]:
+                    if (
+                        "image" in ele
+                        or "image_url" in ele
+                        or "video" in ele
+                        or ele["type"] in ("image", "image_url", "video")
+                    ):
+                        vision_infos.append(ele)
+    return vision_infos
+
+
+def process_vision_info(
+    conversations: list[dict] | list[list[dict]],
+) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None]:
+    vision_infos = extract_vision_info(conversations)
+    # Read images or videos
+    image_inputs = []
+    video_inputs = []
+    for vision_info in vision_infos:
+        if "image" in vision_info or "image_url" in vision_info:
+            image_inputs.append(fetch_image(vision_info))
+        elif "video" in vision_info:
+            video_inputs.append(fetch_video(vision_info))
+        else:
+            raise ValueError("image, image_url or video should in content.")
+    if len(image_inputs) == 0:
+        image_inputs = None
+    if len(video_inputs) == 0:
+        video_inputs = None
+    return image_inputs, video_inputs
@@ -0,0 +1 @@
+ray start --head --port=6399 --num-gpus=1
@@ -0,0 +1 @@
+ray start --address=10.90.67.86:6399
@@ -0,0 +1,8 @@
+python end2end.py --model Qwen/Qwen2.5-Omni-7B \
+                                 --voice-type "m02" \
+                                 --dit-ckpt none \
+                                 --bigvgan-ckpt none \
+                                 --output-wav output_audio \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+                                 --txt-prompts top100.txt
@@ -0,0 +1,11 @@
+export PYTHONPATH=/workspace/omni/vllm-omni/:$PYTHONPATH
+CUDA_VISIBLE_DEVICES=4 python end2end.py --model /workspace/Qwen2.5-Omni-7B/ \
+                                 --voice-type "m02" \
+                                 --dit-ckpt none \
+                                 --bigvgan-ckpt none \
+                                 --output-wav output_audio \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+				 --worker-backend ray \
+				 --ray-address auto \
+                                 --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
@@ -0,0 +1,312 @@
+import tempfile
+from typing import Optional, Union
+from urllib.request import urlopen
+
+import librosa
+import requests
+import resampy
+import soundfile as sf
+import torch
+import torchvision.io
+from processing_omni import fetch_image, fetch_video
+from transformers import AutoConfig, AutoProcessor
+from vllm.inputs import TextPrompt
+
+from vllm_omni.inputs.data import OmniTokensPrompt
+
+# Simple caches to avoid repeated heavy HF loads per prompt
+_PROCESSOR_CACHE: dict[str, "AutoProcessor"] = {}
+_CONFIG_CACHE: dict[str, "AutoConfig"] = {}
+
+
+def get_system_prompt():
+    return {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, "
+                    "Alibaba Group, capable of perceiving auditory and visual inputs, "
+                    "as well as generating text and speech."
+                ),
+            }
+        ],
+    }
+
+
+def resample_wav_to_16khz(input_filepath):
+    data, original_sample_rate = sf.read(input_filepath)
+    # Only use the first channel
+    if len(data.shape) > 1:
+        data = data[:, 0]
+    # resample to 16kHz
+    data_resampled = resampy.resample(data, sr_orig=original_sample_rate, sr_new=16000)
+    return data_resampled
+
+
+def fetch_and_read_video(args, video_url: str, fps=2):
+    def read_video_with_torchvision(video_file_name: str):
+        video, audio, info = torchvision.io.read_video(
+            video_file_name,
+            start_pts=0.0,
+            end_pts=None,
+            pts_unit="sec",
+            output_format="TCHW",
+        )
+
+        total_frames, video_fps = video.size(0), info["video_fps"]
+        total_duration = round(total_frames / video_fps, 3)
+        nframes = int(total_frames / video_fps * fps)
+
+        frame_timestamps = total_duration * torch.arange(1, nframes + 1) / nframes
+        grid_timestamps = frame_timestamps[::2]
+        second_per_grid = grid_timestamps[1] - grid_timestamps[0]
+
+        idx = torch.linspace(0, video.size(0) - 1, nframes).round().long()
+        video = video[idx]
+
+        if args.legacy_omni_video:
+            return [video, total_duration, nframes, second_per_grid.item()]
+        else:
+            return video
+
+    def read_video_with_transformers(video_file_name: Union[str, list[str]]):
+        video, total_duration, nframes, second_per_grid = fetch_video({"video": video_file_name})
+        if total_duration is None and nframes is None:
+            nframes = len(video)
+            total_duration = 0.5 * nframes
+            second_per_grid = 1.0
+        if args.legacy_omni_video:
+            return [video, total_duration, nframes, second_per_grid]
+        else:
+            return video
+
+    def read_video(video_file_name: str):
+        if args.use_torchvision:
+            return read_video_with_torchvision(video_file_name)
+        else:
+            return read_video_with_transformers(video_file_name)
+
+    if isinstance(video_url, str) and video_url.startswith("http"):
+        with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
+            resp = requests.get(video_url)
+            assert resp.status_code == requests.codes.ok, (
+                f"Failed to fetch video from {video_url}, status_code:{resp.status_code}, resp:{resp}"
+            )
+
+            temp_video_file.write(urlopen(video_url).read())
+            temp_video_file_path = temp_video_file.name
+            video_file_name = temp_video_file_path
+            return read_video(video_file_name)
+    else:
+        video_file_name = video_url
+        return read_video(video_file_name)
+
+
+def make_inputs_qwen2_omni(
+    args,
+    messages: list[dict[str, Union[str, list[dict[str, str]]]]],
+    use_audio_in_video: Optional[bool] = False,
+    tokenize: bool = False,
+) -> Union[OmniTokensPrompt, TextPrompt]:
+    from transformers import AutoConfig, AutoProcessor
+
+    # Cached processor/config to prevent per-prompt reloading and repeated warnings
+    if args.model not in _PROCESSOR_CACHE:
+        _PROCESSOR_CACHE[args.model] = AutoProcessor.from_pretrained(args.model)
+    processor = _PROCESSOR_CACHE[args.model]
+
+    config = _CONFIG_CACHE.get(args.model)
+    if config is None:
+        try:
+            config = AutoConfig.from_pretrained(args.model)
+        except Exception:
+            config = None
+        _CONFIG_CACHE[args.model] = config  # cache even if None to avoid retry storms
+
+    # Decide legacy flag only once based on config (default True if unknown)
+    if getattr(args, "legacy_omni_video", None) is None:
+        if config is not None and hasattr(config, "architectures"):
+            args.legacy_omni_video = "Qwen2_5OmniModel" not in config.architectures
+        else:
+            args.legacy_omni_video = True
+
+    audios, images, videos = [], [], []
+    for message in messages:
+        if not isinstance(message["content"], list):
+            message["content"] = [
+                {
+                    "type": "text",
+                    "text": message["content"],
+                }
+            ]
+        index, num_contents = 0, len(message["content"])
+        while index < num_contents:
+            ele = message["content"][index]
+            if "type" not in ele:
+                if "text" in ele:
+                    ele["type"] = "text"
+                elif "audio" in ele:
+                    ele["type"] = "audio"
+                elif "audio_url" in ele:
+                    ele["type"] = "audio_url"
+                elif "image" in ele:
+                    ele["type"] = "image"
+                elif "image_url" in ele:
+                    ele["type"] = "image_url"
+                elif "video" in ele:
+                    ele["type"] = "video"
+                elif "video_url" in ele:
+                    ele["type"] = "video_url"
+                else:
+                    raise ValueError(f"Unknown ele: {ele}")
+
+            if ele["type"] == "audio" or ele["type"] == "audio_url":
+                if "audio_url" in ele:
+                    audio_key = "audio_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_audio_file:
+                        temp_audio_file.write(urlopen(ele[audio_key]).read())
+                        temp_audio_file_path = temp_audio_file.name
+                        audios.append(resample_wav_to_16khz(temp_audio_file_path))
+                        ele["audio"] = temp_audio_file_path
+                elif "audio" in ele:
+                    audio_key = "audio"
+                    audios.append(resample_wav_to_16khz(ele[audio_key]))
+                else:
+                    raise ValueError(f"Unknown ele {ele}")
+            elif use_audio_in_video and (ele["type"] == "video" or ele["type"] == "video_url"):
+                # use video as audio as well
+                if "video_url" in ele:
+                    audio_key = "video_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
+                        temp_video_file.write(urlopen(ele[audio_key]).read())
+                        temp_video_file_path = temp_video_file.name
+                        ele[audio_key] = temp_video_file_path
+                        audios.append(librosa.load(temp_video_file_path, sr=16000)[0])
+                        videos.append(fetch_and_read_video(args, temp_video_file_path))
+                        ele["video"] = temp_video_file_path
+                elif "video" in ele:
+                    audio_key = "video"
+                    audios.append(librosa.load(ele[audio_key], sr=16000)[0])
+                    videos.append(fetch_and_read_video(args, audio_key))
+                else:
+                    raise ValueError(f"Unknown ele {ele}")
+                # insert a audio after the video
+                message["content"].insert(
+                    index + 1,
+                    {
+                        "type": "audio",
+                        "audio": ele[audio_key],
+                    },
+                )
+                # no need to load the added audio again
+                index += 1
+            elif ele["type"] == "video" or ele["type"] == "video_url":
+                if "video_url" in ele:
+                    video_key = "video_url"
+                    with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
+                        temp_video_file.write(urlopen(ele["video_url"]).read())
+                        temp_video_file_path = temp_video_file.name
+                        videos.append(fetch_and_read_video(args, temp_video_file))
+                        ele["video"] = temp_video_file_path
+                else:
+                    video_key = "video"
+                    videos.append(fetch_and_read_video(args, ele[video_key]))
+            elif ele["type"] == "image" or ele["type"] == "image_url":
+                images.append(fetch_image(ele))
+
+            # move to the next content
+            index += 1
+
+    prompt = processor.apply_chat_template(
+        messages,
+        tokenize=tokenize,
+        add_generation_prompt=True,
+        add_vision_id=True,
+    )
+
+    audios = audios if len(audios) > 0 else None
+    images = images if len(images) > 0 else None
+    videos = videos if len(videos) > 0 else None
+
+    multi_modal_data = {}
+    if audios:
+        multi_modal_data["audio"] = audios
+    if images:
+        multi_modal_data["image"] = images
+    if videos:
+        multi_modal_data["video"] = videos
+
+    if isinstance(prompt, list) and isinstance(prompt[0], (list, str)):
+        prompt = prompt[0]
+
+    if tokenize:
+        return OmniTokensPrompt(
+            prompt_token_ids=prompt,
+            multi_modal_data=multi_modal_data,
+        )
+    else:
+        return TextPrompt(
+            prompt=prompt,
+            multi_modal_data=multi_modal_data,
+        )
+
+
+def make_text_prompt(args, prompt):
+    messages = [
+        get_system_prompt(),
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": prompt},
+            ],
+        },
+    ]
+
+    prompt = make_inputs_qwen2_omni(args, messages, tokenize=args.tokenize)
+    return prompt
+
+
+def make_audio_in_video_v2_prompt(args):
+    messages = [
+        {
+            "role": "system",
+            "content": [
+                {
+                    "type": "text",
+                    "text": (
+                        "You are Qwen, a virtual human developed by the Qwen Team, "
+                        "Alibaba Group, capable of perceiving auditory and visual "
+                        "inputs, as well as generating text and speech."
+                    ),
+                }
+            ],
+        },
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "video_url",
+                    "video_url": ("https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-Omni/draw_small.mp4"),
+                },
+            ],
+        },
+    ]
+    prompt = make_inputs_qwen2_omni(
+        args,
+        messages,
+        use_audio_in_video=True,
+        tokenize=args.tokenize,
+    )
+    return prompt
+
+
+def make_omni_prompt(args, prompt=None) -> Union[OmniTokensPrompt, list[OmniTokensPrompt]]:
+    if args.prompt_type == "text":
+        prompt = make_text_prompt(args, prompt)
+    elif args.prompt_type == "audio-in-video-v2":
+        prompt = make_audio_in_video_v2_prompt(args)
+    else:
+        raise ValueError(f"Unsupported prompt type: {args.prompt_type}")
+    return prompt
@@ -6,8 +6,7 @@ Please refer to [README.md](../../../README.md)

 ## Run examples (Qwen2.5-Omni)

-### Launch the Server
-
+Launch the server
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
@@ -17,59 +16,29 @@ If you have custom stage configs file, launch the server with command below
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```

-### Send Multi-modal Request
-
 Get into the example folder
 ```bash
 cd examples/online_serving/qwen2_5_omni
 ```

-####  Send request via python
-
+Send request via python
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
 ```

-The Python client supports the following command-line arguments:
-
- `--query-type` (or `-q`): Query type (default: `mixed_modalities`)
-  - Options: `mixed_modalities`, `use_audio_in_video`, `multi_audios`, `text`
- `--video-path` (or `-v`): Path to local video file or URL
-  - If not provided and query-type uses video, uses default video URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
- `--image-path` (or `-i`): Path to local image file or URL
-  - If not provided and query-type uses image, uses default image URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Supports common image formats: JPEG, PNG, GIF, WebP
-  - Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
- `--audio-path` (or `-a`): Path to local audio file or URL
-  - If not provided and query-type uses audio, uses default audio URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Supports common audio formats: MP3, WAV, OGG, FLAC, M4A
-  - Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
- `--prompt` (or `-p`): Custom text prompt/question
-  - If not provided, uses default prompt for the selected query type
-  - Example: `--prompt "What are the main activities shown in this video?"`
-
-
-For example, to use mixed modalities with all local files:
-
-```bash
-python openai_chat_completion_client_for_multimodal_generation.py \
-    --query-type mixed_modalities \
-    --video-path /path/to/your/video.mp4 \
-    --image-path /path/to/your/image.jpg \
-    --audio-path /path/to/your/audio.wav \
-    --prompt "Analyze all the media content and provide a comprehensive summary."
-```
-
-####  Send request via curl
-
+Send request via curl
 ```bash
 bash run_curl_multimodal_generation.sh mixed_modalities
 ```

+### FAQ
+
+If you encounter error about backend of librosa, try to install ffmpeg with command below.
+```
+sudo apt update
+sudo apt install ffmpeg
+```
+
 ## Run Local Web UI Demo

 This Web UI demo allows users to interact with the model through a web browser.
@@ -84,18 +53,21 @@ python gradio_demo.py  --model Qwen/Qwen2.5-Omni-7B --port 7861

 Then open `http://localhost:7861/` on your local browser to interact with the web UI.

-The gradio script supports the following arguments:

- `--model`: Model name
- `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
- `--port`: Port for Gradio server (default: 7861)
- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
- `--share`: Share the Gradio demo publicly (creates a public link)
+### Options

-### FAQ
+You can customize its basic launch parameters:

-If you encounter error about backend of librosa, try to install ffmpeg with command below.
-```
-sudo apt update
-sudo apt install ffmpeg
+```bash
+python gradio_demo.py \
+    --model Qwen/Qwen2.5-Omni-7B \
+    --ip 127.0.0.1 \
+    --port 7861 \
+    --stage-configs-path /path/to/stage_configs.yaml
 ```
+
+- `--model`: Local model checkpoint to load (default `Qwen/Qwen2.5-Omni-7B`).
+- `--ip`: Host/IP for the Gradio server (default `127.0.0.1`).
+- `--port`: Port for the Gradio server (default `7861`).
+- `--stage-configs-path`: Optional path to custom stage configs YAML.
+- `--share`: Set to expose a temporary public link via Gradio.
@@ -1,6 +1,4 @@
 import base64
-import os
-from typing import Optional

 import requests
 from openai import OpenAI
@@ -30,129 +28,6 @@ def encode_base64_content_from_url(content_url: str) -> str:
    return result


-def encode_base64_content_from_file(file_path: str) -> str:
-    """Encode a local file to base64 format."""
-    with open(file_path, "rb") as f:
-        content = f.read()
-        result = base64.b64encode(content).decode("utf-8")
-    return result
-
-
-def get_video_url_from_path(video_path: Optional[str]) -> str:
-    """Convert a video path (local file or URL) to a video URL format for the API.
-
-    If video_path is None or empty, returns the default URL.
-    If video_path is a local file path, encodes it to base64 data URL.
-    If video_path is a URL, returns it as-is.
-    """
-    if not video_path:
-        # Default video URL
-        return "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
-
-    # Check if it's a URL (starts with http:// or https://)
-    if video_path.startswith(("http://", "https://")):
-        return video_path
-
-    # Otherwise, treat it as a local file path
-    if not os.path.exists(video_path):
-        raise FileNotFoundError(f"Video file not found: {video_path}")
-
-    # Detect video MIME type from file extension
-    video_path_lower = video_path.lower()
-    if video_path_lower.endswith(".mp4"):
-        mime_type = "video/mp4"
-    elif video_path_lower.endswith(".webm"):
-        mime_type = "video/webm"
-    elif video_path_lower.endswith(".mov"):
-        mime_type = "video/quicktime"
-    elif video_path_lower.endswith(".avi"):
-        mime_type = "video/x-msvideo"
-    elif video_path_lower.endswith(".mkv"):
-        mime_type = "video/x-matroska"
-    else:
-        # Default to mp4 if extension is unknown
-        mime_type = "video/mp4"
-
-    video_base64 = encode_base64_content_from_file(video_path)
-    return f"data:{mime_type};base64,{video_base64}"
-
-
-def get_image_url_from_path(image_path: Optional[str]) -> str:
-    """Convert an image path (local file or URL) to an image URL format for the API.
-
-    If image_path is None or empty, returns the default URL.
-    If image_path is a local file path, encodes it to base64 data URL.
-    If image_path is a URL, returns it as-is.
-    """
-    if not image_path:
-        # Default image URL
-        return "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
-
-    # Check if it's a URL (starts with http:// or https://)
-    if image_path.startswith(("http://", "https://")):
-        return image_path
-
-    # Otherwise, treat it as a local file path
-    if not os.path.exists(image_path):
-        raise FileNotFoundError(f"Image file not found: {image_path}")
-
-    # Detect image MIME type from file extension
-    image_path_lower = image_path.lower()
-    if image_path_lower.endswith((".jpg", ".jpeg")):
-        mime_type = "image/jpeg"
-    elif image_path_lower.endswith(".png"):
-        mime_type = "image/png"
-    elif image_path_lower.endswith(".gif"):
-        mime_type = "image/gif"
-    elif image_path_lower.endswith(".webp"):
-        mime_type = "image/webp"
-    else:
-        # Default to jpeg if extension is unknown
-        mime_type = "image/jpeg"
-
-    image_base64 = encode_base64_content_from_file(image_path)
-    return f"data:{mime_type};base64,{image_base64}"
-
-
-def get_audio_url_from_path(audio_path: Optional[str]) -> str:
-    """Convert an audio path (local file or URL) to an audio URL format for the API.
-
-    If audio_path is None or empty, returns the default URL.
-    If audio_path is a local file path, encodes it to base64 data URL.
-    If audio_path is a URL, returns it as-is.
-    """
-    if not audio_path:
-        # Default audio URL
-        return AudioAsset("mary_had_lamb").url
-
-    # Check if it's a URL (starts with http:// or https://)
-    if audio_path.startswith(("http://", "https://")):
-        return audio_path
-
-    # Otherwise, treat it as a local file path
-    if not os.path.exists(audio_path):
-        raise FileNotFoundError(f"Audio file not found: {audio_path}")
-
-    # Detect audio MIME type from file extension
-    audio_path_lower = audio_path.lower()
-    if audio_path_lower.endswith((".mp3", ".mpeg")):
-        mime_type = "audio/mpeg"
-    elif audio_path_lower.endswith(".wav"):
-        mime_type = "audio/wav"
-    elif audio_path_lower.endswith(".ogg"):
-        mime_type = "audio/ogg"
-    elif audio_path_lower.endswith(".flac"):
-        mime_type = "audio/flac"
-    elif audio_path_lower.endswith(".m4a"):
-        mime_type = "audio/mp4"
-    else:
-        # Default to wav if extension is unknown
-        mime_type = "audio/wav"
-
-    audio_base64 = encode_base64_content_from_file(audio_path)
-    return f"data:{mime_type};base64,{audio_base64}"
-
-
 def get_system_prompt():
    return {
        "role": "system",
@@ -169,10 +44,8 @@ def get_system_prompt():
    }


-def get_text_query(custom_prompt: Optional[str] = None):
-    question = (
-        custom_prompt or "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
-    )
+def get_text_query():
+    question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
    prompt = {
        "role": "user",
        "content": [
@@ -185,32 +58,26 @@ def get_text_query(custom_prompt: Optional[str] = None):
    return prompt


-def get_mixed_modalities_query(
-    video_path: Optional[str] = None,
-    image_path: Optional[str] = None,
-    audio_path: Optional[str] = None,
-    custom_prompt: Optional[str] = None,
-):
-    question = (
-        custom_prompt or "What is recited in the audio? What is the content of this image? Why is this video funny?"
-    )
-    video_url = get_video_url_from_path(video_path)
-    image_url = get_image_url_from_path(image_path)
-    audio_url = get_audio_url_from_path(audio_path)
+def get_mixed_modalities_query():
+    question = "What is recited in the audio? What is the content of this image? Why is this video funny?"
    prompt = {
        "role": "user",
        "content": [
            {
                "type": "audio_url",
-                "audio_url": {"url": audio_url},
+                "audio_url": {"url": AudioAsset("mary_had_lamb").url},
            },
            {
                "type": "image_url",
-                "image_url": {"url": image_url},
+                "image_url": {
+                    "url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
+                },
            },
            {
                "type": "video_url",
-                "video_url": {"url": video_url},
+                "video_url": {
+                    "url": "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
+                },
            },
            {
                "type": "text",
@@ -222,9 +89,8 @@ def get_mixed_modalities_query(
    return prompt


-def get_use_audio_in_video_query(video_path: Optional[str] = None, custom_prompt: Optional[str] = None):
-    question = custom_prompt or "Describe the content of the video, then convert what the baby say into text."
-    video_url = get_video_url_from_path(video_path)
+def get_use_audio_in_video_query():
+    question = "Describe the content of the video, then convert what the baby say into text."

    prompt = {
        "role": "user",
@@ -232,7 +98,7 @@ def get_use_audio_in_video_query(video_path: Optional[str] = None, custom_prompt
            {
                "type": "video_url",
                "video_url": {
-                    "url": video_url,
+                    "url": "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4",
                    "num_frames": 16,
                },
            },
@@ -246,15 +112,14 @@ def get_use_audio_in_video_query(video_path: Optional[str] = None, custom_prompt
    return prompt


-def get_multi_audios_query(audio_path: Optional[str] = None, custom_prompt: Optional[str] = None):
-    question = custom_prompt or "Are these two audio clips the same?"
-    audio_url = get_audio_url_from_path(audio_path)
+def get_multi_audios_query():
+    question = "Are these two audio clips the same?"
    prompt = {
        "role": "user",
        "content": [
            {
                "type": "audio_url",
-                "audio_url": {"url": audio_url},
+                "audio_url": {"url": AudioAsset("mary_had_lamb").url},
            },
            {
                "type": "audio_url",
@@ -314,27 +179,7 @@ def run_multimodal_generation(args) -> None:
        code2wav_sampling_params,
    ]

-    # Get paths and custom prompt from args
-    video_path = getattr(args, "video_path", None)
-    image_path = getattr(args, "image_path", None)
-    audio_path = getattr(args, "audio_path", None)
-    custom_prompt = getattr(args, "prompt", None)
-
-    # Get the query function and call it with appropriate parameters
-    query_func = query_map[args.query_type]
-    if args.query_type == "mixed_modalities":
-        prompt = query_func(
-            video_path=video_path, image_path=image_path, audio_path=audio_path, custom_prompt=custom_prompt
-        )
-    elif args.query_type == "use_audio_in_video":
-        prompt = query_func(video_path=video_path, custom_prompt=custom_prompt)
-    elif args.query_type == "multi_audios":
-        prompt = query_func(audio_path=audio_path, custom_prompt=custom_prompt)
-    elif args.query_type == "text":
-        prompt = query_func(custom_prompt=custom_prompt)
-    else:
-        prompt = query_func()
-
+    prompt = query_map[args.query_type]()
    extra_body = {
        "sampling_params_list": sampling_params_list  # Optional, it has a default setting in stage_configs of the corresponding model.
    }
@@ -374,34 +219,6 @@ def parse_args():
        choices=query_map.keys(),
        help="Query type.",
    )
-    parser.add_argument(
-        "--video-path",
-        "-v",
-        type=str,
-        default=None,
-        help="Path to local video file or URL. If not provided and query-type uses video, uses default video URL.",
-    )
-    parser.add_argument(
-        "--image-path",
-        "-i",
-        type=str,
-        default=None,
-        help="Path to local image file or URL. If not provided and query-type uses image, uses default image URL.",
-    )
-    parser.add_argument(
-        "--audio-path",
-        "-a",
-        type=str,
-        default=None,
-        help="Path to local audio file or URL. If not provided and query-type uses audio, uses default audio URL.",
-    )
-    parser.add_argument(
-        "--prompt",
-        "-p",
-        type=str,
-        default=None,
-        help="Custom text prompt/question to use instead of the default prompt for the selected query type.",
-    )

    return parser.parse_args()

@@ -6,8 +6,7 @@ Please refer to [README.md](../../../README.md)

 ## Run examples (Qwen3-Omni)

-### Launch the Server
-
+Launch the server
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
@@ -17,59 +16,21 @@ If you have custom stage configs file, launch the server with command below
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```

-### Send Multi-modal Request
-
 Get into the example folder
 ```bash
 cd examples/online_serving/qwen3_omni
 ```

-####  Send request via python
-
+Send request via python
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
 ```

-The Python client supports the following command-line arguments:
-
- `--query-type` (or `-q`): Query type (default: `use_video`)
-  - Options: `text`, `use_audio`, `use_image`, `use_video`
- `--model` (or `-m`): Model name/path (default: `Qwen/Qwen3-Omni-30B-A3B-Instruct`)
- `--video-path` (or `-v`): Path to local video file or URL
-  - If not provided and query-type is `use_video`, uses default video URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
- `--image-path` (or `-i`): Path to local image file or URL
-  - If not provided and query-type is `use_image`, uses default image URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Supports common image formats: JPEG, PNG, GIF, WebP
-  - Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
- `--audio-path` (or `-a`): Path to local audio file or URL
-  - If not provided and query-type is `use_audio`, uses default audio URL
-  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs
-  - Supports common audio formats: MP3, WAV, OGG, FLAC, M4A
-  - Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
- `--prompt` (or `-p`): Custom text prompt/question
-  - If not provided, uses default prompt for the selected query type
-  - Example: `--prompt "What are the main activities shown in this video?"`
-
-
-For example, to use a local video file with custom prompt:
-
-```bash
-python openai_chat_completion_client_for_multimodal_generation.py \
-    --query-type use_video \
-    --video-path /path/to/your/video.mp4 \
-    --prompt "What are the main activities shown in this video?"
-```
-
-####  Send request via curl
-
+Send request via curl
 ```bash
 bash run_curl_multimodal_generation.sh use_image
 ```

-
 ### FAQ

 If you encounter error about backend of librosa, try to install ffmpeg with command below.
@@ -93,9 +54,23 @@ python gradio_demo.py  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861
 Then open `http://localhost:7861/` on your local browser to interact with the web UI.


-The gradio script supports the following arguments:
+### Options
+
+The gradio demo also supports running with an existing API server and can be customized with the following arguments.
+
+
+```bash
+python gradio_demo.py \
+    --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
+    --use-api-server \
+    --api-base http://localhost:8091/v1 \
+    --ip 127.0.0.1 \
+    --port 7861
+```

 - `--model`: Model name
+- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmni locally.
+- `--api-base`: Base URL for vllm serve (only used when `use-api-server` is set, default: http://localhost:8091/v1)
 - `--ip`: Host/IP for Gradio server (default: 127.0.0.1)
 - `--port`: Port for Gradio server (default: 7861)
 - `--stage-configs-path`: Path to custom stage configs YAML file (optional)
@@ -264,7 +264,7 @@ async def run_inference_async_omni(
        # Process audio
        audio_data = process_audio_file(audio_file)
        if audio_data is not None:
-            multi_modal_data["audio"] = [audio_data] if not isinstance(audio_data, list) else audio_data
+            multi_modal_data["audio"] = audio_data

        # Process image
        image_data = process_image_file(image_file)
@@ -282,9 +282,7 @@ async def run_inference_async_omni(
                video_entry = video_frames
            multi_modal_data["video"] = video_entry
            if use_audio_in_video and extracted_audio is not None and "audio" not in multi_modal_data:
-                multi_modal_data["audio"] = (
-                    [extracted_audio] if not isinstance(extracted_audio, list) else extracted_audio
-                )
+                multi_modal_data["audio"] = extracted_audio
                mm_processor_kwargs["use_audio_in_video"] = True

        # Build the prompt input
@@ -1,6 +1,5 @@
 import base64
-import os
-from typing import NamedTuple, Optional
+from typing import NamedTuple

 import requests
 from openai import OpenAI
@@ -35,129 +34,6 @@ def encode_base64_content_from_url(content_url: str) -> str:
    return result


-def encode_base64_content_from_file(file_path: str) -> str:
-    """Encode a local file to base64 format."""
-    with open(file_path, "rb") as f:
-        content = f.read()
-        result = base64.b64encode(content).decode("utf-8")
-    return result
-
-
-def get_video_url_from_path(video_path: Optional[str]) -> str:
-    """Convert a video path (local file or URL) to a video URL format for the API.
-
-    If video_path is None or empty, returns the default URL.
-    If video_path is a local file path, encodes it to base64 data URL.
-    If video_path is a URL, returns it as-is.
-    """
-    if not video_path:
-        # Default video URL
-        return "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
-
-    # Check if it's a URL (starts with http:// or https://)
-    if video_path.startswith(("http://", "https://")):
-        return video_path
-
-    # Otherwise, treat it as a local file path
-    if not os.path.exists(video_path):
-        raise FileNotFoundError(f"Video file not found: {video_path}")
-
-    # Detect video MIME type from file extension
-    video_path_lower = video_path.lower()
-    if video_path_lower.endswith(".mp4"):
-        mime_type = "video/mp4"
-    elif video_path_lower.endswith(".webm"):
-        mime_type = "video/webm"
-    elif video_path_lower.endswith(".mov"):
-        mime_type = "video/quicktime"
-    elif video_path_lower.endswith(".avi"):
-        mime_type = "video/x-msvideo"
-    elif video_path_lower.endswith(".mkv"):
-        mime_type = "video/x-matroska"
-    else:
-        # Default to mp4 if extension is unknown
-        mime_type = "video/mp4"
-
-    video_base64 = encode_base64_content_from_file(video_path)
-    return f"data:{mime_type};base64,{video_base64}"
-
-
-def get_image_url_from_path(image_path: Optional[str]) -> str:
-    """Convert an image path (local file or URL) to an image URL format for the API.
-
-    If image_path is None or empty, returns the default URL.
-    If image_path is a local file path, encodes it to base64 data URL.
-    If image_path is a URL, returns it as-is.
-    """
-    if not image_path:
-        # Default image URL
-        return "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
-
-    # Check if it's a URL (starts with http:// or https://)
-    if image_path.startswith(("http://", "https://")):
-        return image_path
-
-    # Otherwise, treat it as a local file path
-    if not os.path.exists(image_path):
-        raise FileNotFoundError(f"Image file not found: {image_path}")
-
-    # Detect image MIME type from file extension
-    image_path_lower = image_path.lower()
-    if image_path_lower.endswith((".jpg", ".jpeg")):
-        mime_type = "image/jpeg"
-    elif image_path_lower.endswith(".png"):
-        mime_type = "image/png"
-    elif image_path_lower.endswith(".gif"):
-        mime_type = "image/gif"
-    elif image_path_lower.endswith(".webp"):
-        mime_type = "image/webp"
-    else:
-        # Default to jpeg if extension is unknown
-        mime_type = "image/jpeg"
-
-    image_base64 = encode_base64_content_from_file(image_path)
-    return f"data:{mime_type};base64,{image_base64}"
-
-
-def get_audio_url_from_path(audio_path: Optional[str]) -> str:
-    """Convert an audio path (local file or URL) to an audio URL format for the API.
-
-    If audio_path is None or empty, returns the default URL.
-    If audio_path is a local file path, encodes it to base64 data URL.
-    If audio_path is a URL, returns it as-is.
-    """
-    if not audio_path:
-        # Default audio URL
-        return AudioAsset("mary_had_lamb").url
-
-    # Check if it's a URL (starts with http:// or https://)
-    if audio_path.startswith(("http://", "https://")):
-        return audio_path
-
-    # Otherwise, treat it as a local file path
-    if not os.path.exists(audio_path):
-        raise FileNotFoundError(f"Audio file not found: {audio_path}")
-
-    # Detect audio MIME type from file extension
-    audio_path_lower = audio_path.lower()
-    if audio_path_lower.endswith((".mp3", ".mpeg")):
-        mime_type = "audio/mpeg"
-    elif audio_path_lower.endswith(".wav"):
-        mime_type = "audio/wav"
-    elif audio_path_lower.endswith(".ogg"):
-        mime_type = "audio/ogg"
-    elif audio_path_lower.endswith(".flac"):
-        mime_type = "audio/flac"
-    elif audio_path_lower.endswith(".m4a"):
-        mime_type = "audio/mp4"
-    else:
-        # Default to wav if extension is unknown
-        mime_type = "audio/wav"
-
-    audio_base64 = encode_base64_content_from_file(audio_path)
-    return f"data:{mime_type};base64,{audio_base64}"
-
-
 def get_system_prompt():
    return {
        "role": "system",
@@ -174,10 +50,8 @@ def get_system_prompt():
    }


-def get_text_query(custom_prompt: Optional[str] = None):
-    question = (
-        custom_prompt or "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
-    )
+def get_text_query():
+    question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
    prompt = {
        "role": "user",
        "content": [
@@ -197,15 +71,16 @@ default_system = (
 )


-def get_video_query(video_path: Optional[str] = None, custom_prompt: Optional[str] = None):
-    question = custom_prompt or "Why is this video funny?"
-    video_url = get_video_url_from_path(video_path)
+def get_video_query() -> QueryResult:
+    question = "Why is this video funny?"
    prompt = {
        "role": "user",
        "content": [
            {
                "type": "video_url",
-                "video_url": {"url": video_url},
+                "video_url": {
+                    "url": "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
+                },
            },
            {
                "type": "text",
@@ -216,15 +91,16 @@ def get_video_query(video_path: Optional[str] = None, custom_prompt: Optional[st
    return prompt


-def get_image_query(image_path: Optional[str] = None, custom_prompt: Optional[str] = None):
-    question = custom_prompt or "What is the content of this image?"
-    image_url = get_image_url_from_path(image_path)
+def get_image_query() -> QueryResult:
+    question = "What is the content of this image?"
    prompt = {
        "role": "user",
        "content": [
            {
                "type": "image_url",
-                "image_url": {"url": image_url},
+                "image_url": {
+                    "url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
+                },
            },
            {
                "type": "text",
@@ -235,15 +111,14 @@ def get_image_query(image_path: Optional[str] = None, custom_prompt: Optional[st
    return prompt


-def get_audio_query(audio_path: Optional[str] = None, custom_prompt: Optional[str] = None):
-    question = custom_prompt or "What is the content of this audio?"
-    audio_url = get_audio_url_from_path(audio_path)
+def get_audio_query() -> QueryResult:
+    question = "What is the content of this audio?"
    prompt = {
        "role": "user",
        "content": [
            {
                "type": "audio_url",
-                "audio_url": {"url": audio_url},
+                "audio_url": {"url": AudioAsset("mary_had_lamb").url},
            },
            {
                "type": "text",
@@ -303,25 +178,7 @@ def run_multimodal_generation(args) -> None:
        code2wav_sampling_params,
    ]

-    # Get paths and custom prompt from args
-    video_path = getattr(args, "video_path", None)
-    image_path = getattr(args, "image_path", None)
-    audio_path = getattr(args, "audio_path", None)
-    custom_prompt = getattr(args, "prompt", None)
-
-    # Get the query function and call it with appropriate parameters
-    query_func = query_map[args.query_type]
-    if args.query_type == "use_video":
-        prompt = query_func(video_path=video_path, custom_prompt=custom_prompt)
-    elif args.query_type == "use_image":
-        prompt = query_func(image_path=image_path, custom_prompt=custom_prompt)
-    elif args.query_type == "use_audio":
-        prompt = query_func(audio_path=audio_path, custom_prompt=custom_prompt)
-    elif args.query_type == "text":
-        prompt = query_func(custom_prompt=custom_prompt)
-    else:
-        prompt = query_func()
-
+    prompt = query_map[args.query_type]()
    extra_body = {
        "sampling_params_list": sampling_params_list  # Optional, it has a default setting in stage_configs of the corresponding model.
    }
@@ -368,34 +225,6 @@ def parse_args():
        default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
        help="Model Name / Path",
    )
-    parser.add_argument(
-        "--video-path",
-        "-v",
-        type=str,
-        default=None,
-        help="Path to local video file or URL. If not provided and query-type is 'use_video', uses default video URL.",
-    )
-    parser.add_argument(
-        "--image-path",
-        "-i",
-        type=str,
-        default=None,
-        help="Path to local image file or URL. If not provided and query-type is 'use_image', uses default image URL.",
-    )
-    parser.add_argument(
-        "--audio-path",
-        "-a",
-        type=str,
-        default=None,
-        help="Path to local audio file or URL. If not provided and query-type is 'use_audio', uses default audio URL.",
-    )
-    parser.add_argument(
-        "--prompt",
-        "-p",
-        type=str,
-        default=None,
-        help="Custom text prompt/question to use instead of the default prompt for the selected query type.",
-    )

    return parser.parse_args()

@@ -41,7 +41,7 @@ code2wav_sampling_params='{
  "temperature": 0.0,
  "top_p": 1.0,
  "top_k": -1,
-  "max_tokens": 65536,
+  "max_tokens": 4096 * 16,
  "seed": 42,
  "detokenize": true,
  "repetition_penalty": 1.1
@@ -1,5 +1,5 @@
 site_name: vLLM-Omni
-site_description: Efficient omni-modality model serving for everyone
+site_description: Multi-modality models inference and serving with non-autoregressive structures
 site_author: vLLM-Omni Team
 site_url: https://vllm-project.github.io/vllm-omni/

@@ -5,12 +5,12 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "vllm-omni"
 version = "0.11.0rc1"
-description = "A framework for efficient model inference with omni-modality models"
+description = "vLLM-Omni: Multi-modality models inference and serving with non-autoregressive structures"
 readme = "README.md"
 requires-python = ">=3.9,<3.14"  # Align with vLLM v0.11 supported Python versions
 license = {text = "Apache-2.0"}
 authors = [
-    {name = "vLLM-Omni Team"}
+    {name = "vLLM-Omni Team", email = "vllm-project@gmail.com"}
 ]
 keywords = ["vllm", "multimodal", "diffusion", "transformer", "inference", "serving"]
 classifiers = [
@@ -27,7 +27,6 @@ classifiers = [
    "Topic :: Software Development :: Libraries :: Python Modules",
 ]

-
 dependencies = [
    # Core runtime dependencies (required for actual usage)
    "omegaconf>=2.3.0",
@@ -80,10 +80,6 @@ class OmniDiffusionConfig:
    vae_cpu_offload: bool = True
    pin_cpu_memory: bool = True

-    # VAE memory optimization parameters
-    vae_use_slicing: bool = False
-    vae_use_tiling: bool = False
-
    # STA (Sliding Tile Attention) parameters
    mask_strategy_file_path: str | None = None
    # STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
@@ -24,13 +24,6 @@ def initialize_model(
        module = importlib.import_module(module_name)
        model_class = getattr(module, cls_name)
        model = model_class(od_config=od_config, prefix=mod_relname)
-
-        # Configure VAE memory optimization settings from config
-        if hasattr(model.vae, "use_slicing"):
-            model.vae.use_slicing = od_config.vae_use_slicing
-        if hasattr(model.vae, "use_tiling"):
-            model.vae.use_tiling = od_config.vae_use_tiling
-
        return model
    else:
        raise ValueError(f"Model class {od_config.model_class_name} not found in diffusion model registry.")
@@ -129,7 +129,7 @@ class WorkerProc:
        """
        Receive requests from broadcast queue
        """
-        return self.mq.dequeue(indefinite=True)
+        return self.mq.dequeue()

    # TODO: queueing, cancellation
    def worker_busy_loop(self) -> None:
@@ -5,8 +5,7 @@ import os

 import torch
 import zmq
-from transformers import PretrainedConfig
-from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
+from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm.distributed.parallel_state import (
    init_distributed_environment,
@@ -54,12 +53,7 @@ class NPUWorker:
        torch.npu.set_device(device)

        # hack
-        # set hf_config to a fake one to avolid get attr error
-        class _FakePretrainedConfig(PretrainedConfig):
-            def __getattr__(self, name):
-                return "fake"
-
-        vllm_config = VllmConfig(model_config=ModelConfig(hf_config=_FakePretrainedConfig()))
+        vllm_config = VllmConfig()
        vllm_config.parallel_config.tensor_parallel_size = self.od_config.num_gpus
        set_current_vllm_config(vllm_config)

@@ -135,7 +129,7 @@ class NPUWorkerProc:
        """
        Receive requests from broadcast queue
        """
-        return self.mq.dequeue(indefinite=True)
+        return self.mq.dequeue()

    # TODO: queueing, cancellation
    def worker_busy_loop(self) -> None:
@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from .connectors import (
+    ConnectorSpec, OmniTransferConfig,
+    OmniConnectorBase, InMemoryOmniConnector,
+    OmniConnectorFactory, MooncakeConnector,
+    load_omni_transfer_config, create_simple_config, create_mooncake_config
+)
+
+__all__ = [
+    # Config
+    "ConnectorSpec",
+    "OmniTransferConfig",
+
+    # Connectors
+    "OmniConnectorBase",
+    "InMemoryOmniConnector",
+    "OmniConnectorFactory",
+    "MooncakeConnector",
+
+    # Utilities
+    "load_omni_transfer_config",
+    "create_simple_config",
+    "create_mooncake_config",
+]
@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from .config import ConnectorSpec, OmniTransferConfig
+from .base import OmniConnectorBase, InMemoryOmniConnector
+from .factory import OmniConnectorFactory
+from .mooncake_connector import MooncakeConnector
+from .utils import (
+    load_omni_transfer_config,
+    initialize_connectors_from_config,
+    create_simple_config,
+    create_mooncake_config,
+    get_connectors_config_for_stage,
+    initialize_orchestrator_connectors,
+    get_stage_connector_config,
+    build_stage_connectors,
+)
+
+__all__ = [
+    # Config
+    "ConnectorSpec",
+    "OmniTransferConfig",
+
+    # Base classes and implementations
+    "OmniConnectorBase",
+    "InMemoryOmniConnector",
+
+    # Factory
+    "OmniConnectorFactory",
+
+    # Specific implementations
+    "MooncakeConnector",
+
+    # Utilities
+    "load_omni_transfer_config",
+    "initialize_connectors_from_config",
+    "create_simple_config",
+    "create_mooncake_config",
+    "get_connectors_config_for_stage",
+
+    # Manager helpers
+    "initialize_orchestrator_connectors",
+    "get_stage_connector_config",
+    "build_stage_connectors",
+]
@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+# temporary for compatibility with vllm_omni.entrypoints.omni_stage.py 
+# and vllm_omni.entrypoints.omni_llm.py
+
+import time
+import logging
+from typing import Any, Optional, Dict, Callable
+
+logger = logging.getLogger(__name__)
+
+
+def try_send_via_connector(
+    connector: Any,
+    stage_id: int,
+    next_stage_id: int,
+    req_id: str,
+    next_inputs: Any,
+    sampling_params: Any,
+    original_prompt: Any,
+    next_stage_queue_submit_fn: Callable[[Dict[str, Any]], None],
+    metrics: Any
+) -> bool:
+    """
+    Attempts to send data via OmniConnector.
+    Returns True if successful, False otherwise.
+    Encapsulates the logic of preparing payload, sending via connector, 
+    sending notification, and recording metrics.
+    """
+    try:
+        t0 = time.time()
+
+        # Prepare data for connector
+        payload_data = {
+            "engine_inputs": next_inputs,
+            "sampling_params": sampling_params,
+            "metadata": {
+                "original_prompt": original_prompt,
+                "stage_transition": f"{stage_id}->{next_stage_id}",
+                "timestamp": time.time()
+            }
+        }
+
+        # Send data via connector
+        success, serialized_size = connector.put(str(stage_id), str(next_stage_id), str(req_id), payload_data)
+        
+        if success:
+            # Send lightweight notification via queue
+            notify_payload = {
+                "request_id": req_id,
+                "sampling_params": sampling_params,
+                "from_connector": True,
+                "from_stage": str(stage_id),
+                "to_stage": str(next_stage_id),
+                "sent_ts": time.time(),
+            }
+            next_stage_queue_submit_fn(notify_payload)
+
+            t1 = time.time()
+            tx_ms = (t1 - t0) * 1000.0
+            
+            metrics.on_forward(
+                stage_id,
+                next_stage_id,
+                req_id,
+                serialized_size,  # Use size from connector
+                float(tx_ms),
+                True,  # Mark as using connector
+            )
+            return True
+        else:
+            # If put returned False, we let the caller handle fallback
+            return False
+
+    except Exception as e:
+        logger.warning(
+            "[Orchestrator] OmniConnector failed for req %s: %s; falling back to queue",
+            req_id,
+            e,
+        )
+        return False
+
+
+def try_recv_via_connector(
+    task: Dict[str, Any], 
+    connectors: Dict[Any, Any], 
+    stage_id: int,
+) -> tuple[Any, Optional[Dict[str, Any]]]:
+    """
+    Attempts to resolve input data from either connector or IPC.
+    Returns (engine_inputs, rx_metrics) or (None, None) if failed/skipped.
+    """
+    rid = task["request_id"]
+    
+    if task.get("from_connector"):
+        from_stage = task.get("from_stage")
+        to_stage = str(stage_id)
+
+        if not from_stage:
+            logger.error(
+                "[Stage-%s] 'from_connector' is true but 'from_stage' is missing for request %s", stage_id, rid
+            )
+            return None, None
+
+        # Get connector for this edge
+        connector_key = (from_stage, to_stage)
+        connector = connectors.get(connector_key)
+
+        if connector:
+            try:
+                # Get data from connector with timeout
+                _t_start = time.time()
+                payload_data = connector.get(from_stage, to_stage, str(rid))
+                _t_end = time.time()
+                
+                if payload_data and isinstance(payload_data, dict):
+                    ein = payload_data.get("engine_inputs")
+                    # Use connector's serialization size for consistency
+                    serialized_size = len(connector.serialize_obj(payload_data))
+                    decode_ms = (_t_end - _t_start) * 1000.0
+                    
+                    rx_metrics = {
+                        "rx_decode_time_ms": decode_ms,
+                        "rx_transfer_bytes": serialized_size
+                    }
+                    return ein, rx_metrics
+                else:
+                    logger.error(
+                        "[Stage-%s] Failed to get data from connector for request %s or payload is empty", stage_id, rid
+                    )
+                    return None, None
+            except Exception as e:
+                logger.error(
+                    "[Stage-%s] Error retrieving data from connector for request %s: %s", stage_id, rid, e
+                )
+                return None, None
+        else:
+            logger.error(
+                "[Stage-%s] No connector found for edge %s -> %s for request %s", stage_id, from_stage, to_stage, rid
+            )
+            return None, None
+    else:
+        # Data comes from queue as usual
+        #return maybe_load_from_ipc_with_metrics_fn(task, obj_key="engine_inputs", shm_key="engine_inputs_shm")
+        return None, None
+
@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class OmniConnectorBase(ABC):
+    """Base class for all OmniConnectors."""
+
+    @abstractmethod
+    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int]:
+        """Store Python object, internal serialization handled by connector.
+
+        Args:
+            from_stage: Source stage identifier
+            to_stage: Destination stage identifier
+            request_id: Unique request identifier
+            data: Python object to store
+
+        Returns:
+            tuple: (success: bool, serialized_size: int)
+        """
+        pass
+
+    @abstractmethod
+    def get(self, from_stage: str, to_stage: str, request_id: str) -> Optional[Any]:
+        """Retrieve Python object, internal deserialization handled by connector.
+
+        Args:
+            from_stage: Source stage identifier
+            to_stage: Destination stage identifier
+            request_id: Unique request identifier
+
+        Returns:
+            Python object if found, None otherwise
+        """
+        pass
+
+    @abstractmethod
+    def cleanup(self, request_id: str) -> None:
+        """Clean up resources for a request."""
+        pass
+
+    @abstractmethod
+    def health(self) -> dict[str, Any]:
+        """Return health status and metrics."""
+        pass
+
+    @staticmethod
+    def serialize_obj(obj: Any) -> bytes:
+        """Serialize a Python object to bytes using torch.save (consistent with stage_utils)."""
+        import io
+        import torch
+        buf = io.BytesIO()
+        torch.save(obj, buf)
+        return buf.getvalue()
+
+    @staticmethod
+    def deserialize_obj(data: bytes) -> Any:
+        """Deserialize bytes to Python object using torch.load (consistent with stage_utils)."""
+        import io
+        import torch
+        return torch.load(io.BytesIO(data), map_location="cpu", weights_only=False)
+
+
+class InMemoryOmniConnector(OmniConnectorBase):
+    """Simple in-memory connector for testing and co-located stages."""
+
+    def __init__(self):
+        self._storage: dict[str, Any] = {}
+        self._metrics = {
+            "puts": 0,
+            "gets": 0,
+            "errors": 0,
+        }
+
+    def _make_key(self, from_stage: str, to_stage: str, request_id: str) -> str:
+        return f"{from_stage}:{to_stage}:{request_id}"
+
+    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int]:
+        try:
+            key = self._make_key(from_stage, to_stage, request_id)
+            self._storage[key] = data
+            serialized_size = len(self.serialize_obj(data))
+            self._metrics["puts"] += 1
+            logger.debug(f"InMemoryConnector: stored {key}")
+            return True, serialized_size
+        except Exception as e:
+            self._metrics["errors"] += 1
+            logger.error(f"InMemoryConnector put failed: {e}")
+            return False, 0
+
+    def get(self, from_stage: str, to_stage: str, request_id: str) -> Optional[Any]:
+        try:
+            key = self._make_key(from_stage, to_stage, request_id)
+            data = self._storage.get(key)
+            if data is not None:
+                self._metrics["gets"] += 1
+                return data
+            return None
+        except Exception as e:
+            logger.error(f"InMemoryConnector get failed: {e}")
+            return None
+
+    def cleanup(self, request_id: str) -> None:
+        keys_to_remove = [k for k in self._storage.keys() if k.endswith(f":{request_id}")]
+        for key in keys_to_remove:
+            del self._storage[key]
+        logger.debug(f"InMemoryConnector: cleaned up {len(keys_to_remove)} keys for {request_id}")
+
+    def health(self) -> dict[str, Any]:
+        return {
+            "status": "healthy",
+            "storage_size": len(self._storage),
+            **self._metrics,
+        }
@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ConnectorSpec:
+    """Specification for a connector instance."""
+    name: str  # e.g., "MooncakeConnector", "InMemoryConnector"
+    extra: dict[str, Any] = field(default_factory=dict)  # backend-specific config
+
+
+@dataclass
+class OmniTransferConfig:
+    """Top-level configuration for OmniConnector system."""
+    # Direct mapping: (from_stage, to_stage) -> connector
+    connectors: dict[tuple[str, str], ConnectorSpec] = field(default_factory=dict)
+    default_connector: Optional[ConnectorSpec] = None
+
+    def get_connector_for_edge(self, from_stage: str, to_stage: str) -> Optional[ConnectorSpec]:
+        """Get connector spec for a specific edge."""
+        edge_key = (from_stage, to_stage)
+        return self.connectors.get(edge_key, self.default_connector)
+
+    def has_connector_for_edge(self, from_stage: str, to_stage: str) -> bool:
+        """Check if there's a connector configured for the edge."""
+        return self.get_connector_for_edge(from_stage, to_stage) is not None
@@ -0,0 +1,9 @@
+export PYTHONPATH=/workspace/omni/vllm-omni/:$PYTHONPATH
+CUDA_VISIBLE_DEVICES=0,1,2,3 python end2end.py --model /workspace/Qwen2.5-Omni-7B/ \
+                                 --voice-type "m02" \
+                                 --dit-ckpt none \
+                                 --bigvgan-ckpt none \
+                                 --output-wav output_audio \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+                                 --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+"""
+Test OmniConnector integration with OmniLLM.
+"""
+
+import sys
+import os
+sys.path.insert(0, os.path.dirname(__file__))
+
+from vllm_omni.distributed.connectors import create_mooncake_config
+
+def test_connector_config():
+    """Test connector configuration creation."""
+    print("Testing connector configuration...")
+
+    # Create Mooncake config
+    config = create_mooncake_config(
+        host="127.0.0.1",
+        metadata_server="http://127.0.0.1:8080/metadata",
+        master="127.0.0.1:50051"
+    )
+
+    print(f"Created config with {len(config.connectors)} connectors")
+    for edge, connector_spec in config.connectors.items():
+        print(f"  {edge}: {connector_spec.name}")
+
+    return config
+
+def test_omni_llm_initialization():
+    """Test OmniLLM with connector config."""
+    print("\nTesting OmniLLM initialization with connectors...")
+
+    try:
+        from vllm_omni.entrypoints.omni_llm import OmniLLM
+
+        # Create connector config
+        connector_config = {
+            "connectors": {
+                "0_to_1": {
+                    "name": "InMemoryConnector",  # Use InMemory for testing
+                    "extra": {}
+                }
+            }
+        }
+
+        # Test connector initialization logic
+        omni = OmniLLM.__new__(OmniLLM)  # Create without calling __init__
+        omni.omni_transfer_config = connector_config
+        omni.connectors = {}
+        omni._initialize_connectors()
+
+        # Check if connector was created
+        if ("0", "1") in omni.connectors:
+            print(f"✅ Connector created: {type(omni.connectors[('0', '1')]).__name__}")
+            return True
+        else:
+            print("❌ Connector not created")
+            return False
+
+    except Exception as e:
+        print(f"Error: {e}")
+        return False
+
+if __name__ == "__main__":
+    print("OmniConnector Integration Test")
+    print("=" * 40)
+
+    # Test config creation
+    config = test_connector_config()
+
+    # Test OmniLLM integration
+    success = test_omni_llm_initialization()
+
+    if success:
+        print("\n✅ Integration test passed!")
+        print("Connector system is ready for use with OmniLLM")
+    else:
+        print("\n❌ Integration test failed!")
@@ -0,0 +1,332 @@
+#!/usr/bin/env python3
+"""
+Test EPD (Encode-Prefill-Decode) 3-stage pipeline with connectors.
+Tests data flow between stages using configured connectors.
+"""
+
+import os
+import sys
+from typing import Any, Dict, Optional
+from dataclasses import dataclass
+import uuid
+
+try:
+    import yaml
+except ImportError:
+    print("PyYAML not available, using basic YAML parsing")
+    yaml = None
+
+# Add the connectors package to path
+current_dir = os.path.dirname(os.path.abspath(__file__))
+parent_dir = os.path.dirname(current_dir)
+sys.path.insert(0, parent_dir)
+
+# Import from parent package
+try:
+    from connectors.config import ConnectorSpec, OmniTransferConfig
+    from connectors.base import OmniConnectorBase, InMemoryOmniConnector
+    from connectors.factory import OmniConnectorFactory
+    from connectors.utils import load_omni_transfer_config
+except ImportError:
+    # Fallback for direct execution
+    sys.path.insert(0, current_dir)
+    from config import ConnectorSpec, OmniTransferConfig
+    from base import OmniConnectorBase, InMemoryOmniConnector
+    from factory import OmniConnectorFactory
+    from utils import load_omni_transfer_config
+
+
+@dataclass
+class StageResult:
+    """Result from a stage processing."""
+    stage_name: str
+    request_id: str
+    output_data: Any
+    metadata: Dict[str, Any]
+
+
+class EncodeStage:
+    """Simple Encode stage for testing."""
+
+    def __init__(self, stage_id: int, output_connector: Optional[OmniConnectorBase] = None):
+        self.stage_id = stage_id
+        self.stage_name = "encode"
+        self.output_connector = output_connector
+
+    def process(self, request_id: str, input_data: Any) -> StageResult:
+        """Process input data and optionally send to next stage."""
+        print(f"[{self.stage_name}] Processing request {request_id}")
+
+        # Simple encoding simulation
+        encoded_data = {
+            "original_text": input_data,
+            "encoded_tokens": [ord(c) for c in str(input_data)][:10],  # Simple tokenization
+            "embeddings": [0.1 * i for i in range(10)],  # Mock embeddings
+            "stage": self.stage_name
+        }
+
+        result = StageResult(
+            stage_name=self.stage_name,
+            request_id=request_id,
+            output_data=encoded_data,
+            metadata={"tokens_count": len(encoded_data["encoded_tokens"])}
+        )
+
+        # Send to next stage if connector is available
+        if self.output_connector:
+            success = self.output_connector.put(
+                str(self.stage_id), "1", request_id, encoded_data
+            )
+            print(f"[{self.stage_name}] Sent data to stage 1 via connector: {success}")
+
+        return result
+
+
+class PrefillStage:
+    """Simple Prefill stage for testing."""
+
+    def __init__(self, stage_id: int,
+                 input_connector: Optional[OmniConnectorBase] = None,
+                 output_connector: Optional[OmniConnectorBase] = None):
+        self.stage_id = stage_id
+        self.stage_name = "prefill"
+        self.input_connector = input_connector
+        self.output_connector = output_connector
+
+    def process(self, request_id: str) -> Optional[StageResult]:
+        """Process data from previous stage."""
+        print(f"[{self.stage_name}] Processing request {request_id}")
+
+        # Try to get data from input connector
+        input_data = None
+        if self.input_connector:
+            input_data = self.input_connector.get("0", str(self.stage_id), request_id)
+
+        if input_data is None:
+            print(f"[{self.stage_name}] No input data received for {request_id}")
+            return None
+
+        print(f"[{self.stage_name}] Received data from encode")
+
+        # Simple prefill simulation
+        prefills_data = {
+            **input_data,
+            "kv_cache": {"layer_0": [0.2 * i for i in range(20)]},  # Mock KV cache
+            "stage": self.stage_name
+        }
+
+        result = StageResult(
+            stage_name=self.stage_name,
+            request_id=request_id,
+            output_data=prefills_data,
+            metadata={"kv_size": len(prefills_data["kv_cache"])}
+        )
+
+        # Send to next stage if connector is available
+        if self.output_connector:
+            success = self.output_connector.put(
+                str(self.stage_id), "2", request_id, prefills_data
+            )
+            print(f"[{self.stage_name}] Sent data to stage 2 via connector: {success}")
+
+        return result
+
+
+class DecodeStage:
+    """Simple Decode stage for testing (final stage in 3-stage pipeline)."""
+
+    def __init__(self, stage_id: int,
+                 input_connector: Optional[OmniConnectorBase] = None):
+        self.stage_id = stage_id
+        self.stage_name = "decode"
+        self.input_connector = input_connector
+
+    def process(self, request_id: str) -> Optional[StageResult]:
+        """Process data from previous stage (final output)."""
+        print(f"[{self.stage_name}] Processing request {request_id}")
+
+        # Try to get data from input connector
+        input_data = None
+        if self.input_connector:
+            input_data = self.input_connector.get("1", str(self.stage_id), request_id)
+
+        if input_data is None:
+            print(f"[{self.stage_name}] No input data received for {request_id}")
+            return None
+
+        print(f"[{self.stage_name}] Received data from prefill")
+
+        # Simple decode simulation (final output)
+        decoded_data = {
+            **input_data,
+            "hidden_states": [0.3 * i for i in range(15)],  # Mock hidden states
+            "final_output": f"Decoded text from: {input_data.get('original_text', 'unknown')}",
+            "stage": self.stage_name
+        }
+
+        result = StageResult(
+            stage_name=self.stage_name,
+            request_id=request_id,
+            output_data=decoded_data,
+            metadata={"hidden_size": len(decoded_data["hidden_states"]), "final": True}
+        )
+
+        return result
+
+
+class GeneratorStage:
+    """Simple Generator stage for testing (removed for 3-stage pipeline)."""
+    pass  # Not used in current 3-stage setup
+
+
+def load_config_from_yaml(yaml_path: str) -> OmniTransferConfig:
+    """Load connector configuration from YAML file."""
+    if yaml is None:
+        # Fallback: create a simple test config
+        print("Using fallback InMemoryConnector config for testing")
+        return create_test_config()
+
+    with open(yaml_path, 'r', encoding='utf-8') as f:
+        config_dict = yaml.safe_load(f)
+
+    try:
+        config = load_omni_transfer_config(config_dict=config_dict)
+        print(f"Successfully loaded config from YAML with {len(config.connectors)} connectors")
+        return config
+    except Exception as e:
+        print(f"Failed to load YAML config: {e}, using fallback config")
+        return create_test_config()
+
+
+def create_test_config() -> OmniTransferConfig:
+    """Create a test configuration with InMemoryConnectors."""
+    from .config import ConnectorSpec
+
+    connectors = {}
+    # Create connectors for EPD (3-stage) pipeline: 0->1->2
+    pipeline_stages = [("0", "1"), ("1", "2")]
+
+    for from_stage, to_stage in pipeline_stages:
+        connectors[(from_stage, to_stage)] = ConnectorSpec(
+            name="InMemoryConnector",
+            extra={}
+        )
+
+    return OmniTransferConfig(connectors=connectors)
+
+
+def create_stages_with_connectors(config: OmniTransferConfig) -> Dict[str, Any]:
+    """Create EPD (3-stage) pipeline with appropriate connectors."""
+
+    # Create connector instances based on config
+    connectors = {}
+    for edge_key, connector_spec in config.connectors.items():
+        from_stage, to_stage = edge_key
+        connectors[edge_key] = OmniConnectorFactory.create_connector(connector_spec)
+
+    # Create stages with connectors (3-stage pipeline: 0->1->2)
+    stages = {}
+
+    # Encode stage (stage 0)
+    encode_output_connector = connectors.get(("0", "1"))  # to prefill
+    stages["encode"] = EncodeStage(0, encode_output_connector)
+
+    # Prefill stage (stage 1)
+    prefill_input_connector = connectors.get(("0", "1"))   # from encode
+    prefill_output_connector = connectors.get(("1", "2"))  # to decode
+    stages["prefill"] = PrefillStage(1, prefill_input_connector, prefill_output_connector)
+
+    # Decode stage (stage 2) - final stage
+    decode_input_connector = connectors.get(("1", "2"))    # from prefill
+    stages["decode"] = DecodeStage(2, decode_input_connector)
+
+    return stages
+
+
+def test_pipeline_with_config(yaml_path: str):
+    """Test the EPDG pipeline with connector configuration."""
+
+    print("🔧 Loading configuration from YAML...")
+    try:
+        config = load_config_from_yaml(yaml_path)
+        print(f"✓ Loaded config with {len(config.connectors)} connector configurations")
+    except Exception as e:
+        print(f"✗ Failed to load config: {e}")
+        return False
+
+    print("\n🏗️  Creating EPDG stages with connectors...")
+    try:
+        stages = create_stages_with_connectors(config)
+        print("✓ Created all stages with connectors")
+    except Exception as e:
+        print(f"✗ Failed to create stages: {e}")
+        return False
+
+    print("\n🚀 Testing data flow through pipeline...")
+
+    # Test request
+    request_id = f"test_req_{uuid.uuid4()}"
+    input_text = "Hello, EPDG pipeline!"
+
+    try:
+        # Stage 1: Encode
+        print(f"\n--- Stage 1: Encode ---")
+        encode_result = stages["encode"].process(request_id, input_text)
+        print(f"✓ Encode completed: {encode_result.metadata}")
+
+        # Stage 2: Prefill
+        print(f"\n--- Stage 2: Prefill ---")
+        prefill_result = stages["prefill"].process(request_id)
+        if prefill_result:
+            print(f"✓ Prefill completed: {prefill_result.metadata}")
+        else:
+            print("✗ Prefill failed - no input data")
+            return False
+
+        # Stage 3: Decode
+        print(f"\n--- Stage 3: Decode ---")
+        decode_result = stages["decode"].process(request_id)
+        if decode_result:
+            print(f"✓ Decode completed: {decode_result.metadata}")
+        else:
+            print("✗ Decode failed - no input data")
+            return False
+
+        # Decode stage produces final output in 3-stage pipeline
+
+    except Exception as e:
+        print(f"✗ Pipeline test failed: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+    print("\n🎉 Pipeline test completed successfully!")
+    return True
+
+
+def main():
+    """Main test function."""
+    # Path to the YAML config file
+    yaml_path = "/workspace/omni/vllm-omni/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml"
+
+    if not os.path.exists(yaml_path):
+        print(f"✗ YAML config file not found: {yaml_path}")
+        return False
+
+    print("Testing EPDG Pipeline with Connectors")
+    print("=" * 50)
+    print(f"Config file: {yaml_path}")
+
+    success = test_pipeline_with_config(yaml_path)
+
+    if success:
+        print("\n✅ All tests passed!")
+    else:
+        print("\n❌ Tests failed!")
+
+    return success
+
+
+if __name__ == "__main__":
+    success = main()
+    sys.exit(0 if success else 1)
@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+try:
+    from .base import OmniConnectorBase
+    from .config import ConnectorSpec
+except ImportError:
+    # Fallback for direct execution
+    import sys
+    import os
+    sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+    from connectors.base import OmniConnectorBase
+    from connectors.config import ConnectorSpec
+
+logger = logging.getLogger(__name__)
+
+
+class OmniConnectorFactory:
+    """Factory for creating OmniConnectors."""
+
+    _registry: dict[str, Callable[[dict[str, Any]], OmniConnectorBase]] = {}
+
+    @classmethod
+    def register_connector(
+        cls,
+        name: str,
+        constructor: Callable[[dict[str, Any]], OmniConnectorBase]
+    ) -> None:
+        """Register a connector constructor."""
+        if name in cls._registry:
+            raise ValueError(f"Connector '{name}' is already registered.")
+        cls._registry[name] = constructor
+        logger.debug(f"Registered connector: {name}")
+
+    @classmethod
+    def create_connector(cls, spec: ConnectorSpec) -> OmniConnectorBase:
+        """Create a connector from specification."""
+        if spec.name not in cls._registry:
+            raise ValueError(f"Unknown connector: {spec.name}. Available: {list(cls._registry.keys())}")
+
+        constructor = cls._registry[spec.name]
+        try:
+            connector = constructor(spec.extra)
+            logger.info(f"Created connector: {spec.name}")
+            return connector
+        except Exception as e:
+            logger.error(f"Failed to create connector {spec.name}: {e}")
+            raise
+
+    @classmethod
+    def list_registered_connectors(cls) -> list[str]:
+        """List all registered connector names."""
+        return list(cls._registry.keys())
+
+
+# Register built-in connectors with lazy imports
+def _create_mooncake_connector(config: dict[str, Any]) -> OmniConnectorBase:
+    try:
+        from .mooncake_connector import MooncakeConnector
+    except ImportError:
+        # Fallback import
+        import sys
+        import os
+        sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+        from connectors.mooncake_connector import MooncakeConnector
+    return MooncakeConnector(config)
+
+
+def _create_in_memory_connector(config: dict[str, Any]) -> OmniConnectorBase:
+    try:
+        from .base import InMemoryOmniConnector
+    except ImportError:
+        # Fallback import
+        import sys
+        import os
+        sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+        from connectors.base import InMemoryOmniConnector
+    return InMemoryOmniConnector()
+
+
+# Register connectors
+OmniConnectorFactory.register_connector("MooncakeConnector", _create_mooncake_connector)
+OmniConnectorFactory.register_connector("InMemoryConnector", _create_in_memory_connector)
@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import io
+import logging
+import time
+from typing import Any, Optional
+
+import torch
+
+from .base import OmniConnectorBase
+
+logger = logging.getLogger(__name__)
+
+try:
+    from mooncake.store import MooncakeDistributedStore, ReplicateConfig
+except ImportError:
+    try:
+        from mooncake import MooncakeDistributedStore, ReplicateConfig
+    except ImportError:
+        logger.warning("Mooncake not available, MooncakeOmniConnector will not work")
+        MooncakeDistributedStore = None
+        ReplicateConfig = None
+
+
+class MooncakeConnector(OmniConnectorBase):
+    """Mooncake-based distributed connector for OmniConnector."""
+
+    def __init__(self, config: dict[str, Any]):
+        if MooncakeDistributedStore is None:
+            raise ImportError("Mooncake not available")
+
+        self.config = config
+        self.host = config.get("host", "127.0.0.1")
+        self.metadata = config.get("metadata_server", "http://127.0.0.1:8080/metadata")
+        self.master = config.get("master", "127.0.0.1:50051")
+        self.segment = config.get("segment", 512 * 1024 * 1024)  # 512MB
+        self.localbuf = config.get("localbuf", 64 * 1024 * 1024)  # 64MB
+        self.proto = config.get("proto", "tcp")
+        self.rdma = config.get("rdma", "")
+
+        self.store: Optional[MooncakeDistributedStore] = None
+        self.pin: Optional[ReplicateConfig] = None
+
+        self._metrics = {
+            "puts": 0,
+            "gets": 0,
+            "bytes_transferred": 0,
+            "errors": 0,
+            "timeouts": 0,
+        }
+
+        self._init_store()
+
+    def _make_key(self, rid: str, from_stage: str, to_stage: str) -> str:
+        """Generate store key for request between stages."""
+        return f"{rid}/{from_stage}_to_{to_stage}"
+
+    def _init_store(self):
+        """Initialize Mooncake store."""
+        try:
+            self.store = MooncakeDistributedStore()
+            rc = self.store.setup(
+                self.host, self.metadata, self.segment,
+                self.localbuf, self.proto, self.rdma, self.master
+            )
+            if rc != 0:
+                raise RuntimeError(f"Mooncake setup failed: {rc}")
+
+            self.pin = ReplicateConfig()
+            self.pin.with_soft_pin = True
+            logger.info("MooncakeConnector initialized successfully")
+        except Exception as e:
+            logger.error(f"Failed to initialize Mooncake store: {e}")
+            raise
+
+    # Use base class serialization methods for consistency
+
+    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int]:
+        if not self.store:
+            logger.error("Store not initialized")
+            return False, 0
+
+        try:
+            serialized_data = self.serialize_obj(data)
+            key = self._make_key(request_id, from_stage, to_stage)
+            self.store.put(key, serialized_data, self.pin)
+
+            self._metrics["puts"] += 1
+            self._metrics["bytes_transferred"] += len(serialized_data)
+
+            logger.debug(f"MooncakeConnector: stored {key}, {len(serialized_data)} bytes")
+            print(f"========= mooncake put {from_stage}-{to_stage}-{request_id}, with size {len(serialized_data)}")
+            return True, len(serialized_data)
+
+        except Exception as e:
+            self._metrics["errors"] += 1
+            logger.error(f"MooncakeConnector put failed: {e}")
+            return False, 0
+
+    def get(self, from_stage: str, to_stage: str, request_id: str) -> Optional[Any]:
+        if not self.store:
+            logger.error("Store not initialized")
+            return None
+
+        retries = 20
+        sleep_s = 0.05
+        key = self._make_key(request_id, from_stage, to_stage)
+
+        for attempt in range(retries):
+            try:
+                raw_data = self.store.get(key)
+
+                if raw_data:
+                    data = self.deserialize_obj(raw_data)
+                    self._metrics["gets"] += 1
+                    print(f"========= mooncake get {from_stage}-{to_stage}-{request_id}, with size {len(raw_data)}")
+                    return data
+
+            except Exception as e:
+                logger.debug(f"MooncakeConnector get attempt {attempt} failed: {e}")
+
+            if attempt < retries - 1:
+                time.sleep(sleep_s)
+
+        self._metrics["timeouts"] += 1
+        logger.warning(f"MooncakeConnector: timeout waiting for {key}")
+        return None
+
+    def cleanup(self, request_id: str) -> None:
+        if not self.store:
+            return
+
+        # Note: Mooncake doesn't have explicit delete, data will be garbage collected
+        # We could implement a cleanup mechanism by storing deletion markers
+        logger.debug(f"MooncakeConnector: cleanup requested for {request_id} (no-op)")
+
+    def health(self) -> dict[str, Any]:
+        if not self.store:
+            return {"status": "unhealthy", "error": "Store not initialized"}
+
+        return {
+            "status": "healthy",
+            "host": self.host,
+            "metadata_server": self.metadata,
+            "master": self.master,
+            **self._metrics,
+        }
+
+    def close(self):
+        """Clean shutdown."""
+        if self.store:
+            try:
+                self.store.close()
+                self.store = None
+                logger.info("MooncakeConnector closed")
+            except Exception as e:
+                logger.error(f"Error closing Mooncake store: {e}")
@@ -0,0 +1,314 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Utilities for OmniConnector configuration and validation."""
+
+import json
+import logging
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple, Union
+
+from vllm_omni.distributed.connectors.config import OmniTransferConfig, ConnectorSpec
+from vllm_omni.distributed.connectors.factory import OmniConnectorFactory
+from vllm_omni.distributed.connectors.base import OmniConnectorBase
+
+logger = logging.getLogger(__name__)
+
+
+def initialize_connectors_from_config(
+    config_path: Optional[Union[str, Path]] = None
+) -> tuple[Optional[OmniTransferConfig], dict[tuple[str, str], OmniConnectorBase]]:
+    """
+    Initialize connectors from configuration file.
+
+    Returns:
+        tuple: (OmniTransferConfig, dict of {(from, to): connector_instance})
+    """
+    try:
+        transfer_config = load_omni_transfer_config(config_path)
+    except Exception as e:
+        logger.warning(f"Failed to load OmniTransferConfig from {config_path}: {e}")
+        return None, {}
+
+    if not transfer_config:
+        logger.info("No OmniTransferConfig provided, using default IPC")
+        return None, {}
+
+    # 使用统一的连接器创建逻辑
+    connectors = create_connectors_from_config(transfer_config.connectors)
+    return transfer_config, connectors
+
+
+def create_connectors_from_config(
+    connectors_config: dict[tuple[str, str], ConnectorSpec]
+) -> dict[tuple[str, str], OmniConnectorBase]:
+    """通用连接器创建逻辑，从配置字典创建连接器实例。"""
+    connectors = {}
+    try:
+        for edge_key, connector_spec in connectors_config.items():
+            connector = OmniConnectorFactory.create_connector(connector_spec)
+            connectors[edge_key] = connector
+            logger.info(f"Created connector for {edge_key[0]} -> {edge_key[1]}: {type(connector).__name__}")
+    except Exception as e:
+        logger.error(f"Failed to initialize connectors: {e}")
+
+    return connectors
+
+
+def get_connectors_config_for_stage(
+    transfer_config: Optional[OmniTransferConfig],
+    stage_id: Union[str, int]
+) -> dict[str, Any]:
+    """
+    Extract connector configurations relevant for a specific stage worker.
+    
+    Returns a dict compatible with worker initialization:
+    {
+        "from_stage_X": {
+            "spec": {
+                "name": "ConnectorName",
+                "extra": {...}
+            }
+        },
+        ...
+    }
+    """
+    if not transfer_config:
+        return {}
+        
+    stage_connectors_config = {}
+    target_stage = str(stage_id)
+    
+    # Iterate through all configured edges
+    for (from_stage, to_stage), spec in transfer_config.connectors.items():
+        # We only care about incoming edges for the worker process
+        # (Worker needs to create connectors to receive data)
+        if to_stage == target_stage:
+            stage_connectors_config[f"from_stage_{from_stage}"] = {
+                "spec": {
+                    "name": spec.name,
+                    "extra": spec.extra
+                }
+            }
+            
+    return stage_connectors_config
+
+
+def load_omni_transfer_config(
+    config_path: Optional[Union[str, Path]] = None,
+    config_dict: Optional[dict[str, Any]] = None
+) -> Optional[OmniTransferConfig]:
+    """Load OmniTransferConfig from file or dict."""
+    if config_path is None and config_dict is None:
+        return None
+
+    if config_path is not None:
+        config_path = Path(config_path)
+        if not config_path.exists():
+            raise FileNotFoundError(f"Config file not found: {config_path}")
+
+        with open(config_path, 'r', encoding='utf-8') as f:
+            if config_path.suffix.lower() == '.json':
+                config_dict = json.load(f)
+            elif config_path.suffix.lower() in ['.yaml', '.yml']:
+                try:
+                    import yaml
+                    config_dict = yaml.safe_load(f)
+                except ImportError:
+                    raise ImportError("PyYAML required for YAML config files")
+            else:
+                raise ValueError(f"Unsupported config file format: {config_path.suffix}")
+
+    if config_dict is None:
+        return None
+
+    # Parse connectors
+    connectors = {}
+    runtime_config = config_dict.get('runtime', {})
+
+    # Parse global connectors (from runtime.connectors)
+    global_connectors = runtime_config.get('connectors', {})
+    for conn_name, conn_config in global_connectors.items():
+        connector = ConnectorSpec(
+            name=conn_config['name'],
+            extra=conn_config.get('extra', {})
+        )
+        # Store globally defined connectors for reference
+        # These will be resolved to specific edges later
+
+    # Parse stage-level connectors
+    for stage_config in config_dict.get('stage_args', []):
+        stage_id = str(stage_config['stage_id'])
+
+        # Input connectors
+        for input_key, conn_ref in stage_config.get('input_connectors', {}).items():
+            if isinstance(conn_ref, str):
+                # Reference to global connector
+                if conn_ref in global_connectors:
+                    conn_config = global_connectors[conn_ref]
+                    connector = ConnectorSpec(
+                        name=conn_config['name'],
+                        extra=conn_config.get('extra', {})
+                    )
+                else:
+                    raise ValueError(f"Undefined connector reference: {conn_ref}")
+            else:
+                # Inline connector definition
+                connector = ConnectorSpec(
+                    name=conn_ref['name'],
+                    extra=conn_ref.get('extra', {})
+                )
+
+            # Parse from_stage from key (e.g., "from_stage_0" -> "0")
+            from_stage = input_key.replace("from_stage_", "")
+            edge_key = (from_stage, stage_id)
+            connectors[edge_key] = connector
+
+        # Output connectors
+        for output_key, conn_ref in stage_config.get('output_connectors', {}).items():
+            if isinstance(conn_ref, str):
+                # Reference to global connector
+                if conn_ref in global_connectors:
+                    conn_config = global_connectors[conn_ref]
+                    connector = ConnectorSpec(
+                        name=conn_config['name'],
+                        extra=conn_config.get('extra', {})
+                    )
+                else:
+                    raise ValueError(f"Undefined connector reference: {conn_ref}")
+            else:
+                # Inline connector definition
+                connector = ConnectorSpec(
+                    name=conn_ref['name'],
+                    extra=conn_ref.get('extra', {})
+                )
+
+            # Parse to_stage from key (e.g., "to_stage_1" -> "1")
+            to_stage = output_key.replace("to_stage_", "")
+            edge_key = (stage_id, to_stage)
+            connectors[edge_key] = connector
+
+    config = OmniTransferConfig(connectors=connectors)
+
+    logger.info(f"Loaded OmniTransferConfig with {len(connectors)} connector configurations")
+    return config
+
+
+def create_simple_config(
+    connector_name: str = "InMemoryConnector",
+    connector_config: Optional[dict[str, Any]] = None
+) -> OmniTransferConfig:
+    """Create a simple OmniTransferConfig for testing."""
+    if connector_config is None:
+        connector_config = {}
+
+    # Create connectors for typical EPDG pipeline
+    connectors = {}
+    pipeline_edges = [
+        ("0", "1"),  # thinker -> talker
+        ("1", "2"),  # talker -> code2wav
+    ]
+
+    for from_stage, to_stage in pipeline_edges:
+        connectors[(from_stage, to_stage)] = ConnectorSpec(
+            name=connector_name,
+            extra=connector_config
+        )
+
+    return OmniTransferConfig(connectors=connectors)
+
+
+def create_mooncake_config(
+    host: str = "127.0.0.1",
+    metadata_server: str = "http://127.0.0.1:8080/metadata",
+    master: str = "127.0.0.1:50051",
+    segment: int = 512 * 1024 * 1024,
+    localbuf: int = 64 * 1024 * 1024,
+) -> OmniTransferConfig:
+    """Create OmniTransferConfig for Mooncake-based setup."""
+    connector_config = {
+        "host": host,
+        "metadata_server": metadata_server,
+        "master": master,
+        "segment": segment,
+        "localbuf": localbuf,
+    }
+
+    return create_simple_config("MooncakeConnector", connector_config)
+
+
+# High-level management functions (moved from manager.py)
+
+def initialize_orchestrator_connectors(
+    config_path: Optional[str],
+) -> Tuple[Optional[OmniTransferConfig], Dict[tuple[str, str], OmniConnectorBase]]:
+    """Initialize connectors shared at orchestrator level."""
+    try:
+        transfer_config, connectors = initialize_connectors_from_config(config_path)
+        return transfer_config, connectors
+    except Exception as exc:  # pragma: no cover - defensive logging
+        logger.error("Unexpected error initializing connectors: %s", exc)
+        return None, {}
+
+
+def get_stage_connector_config(
+    transfer_config: Optional[OmniTransferConfig],
+    stage_id: int,
+) -> dict[str, Any]:
+    """Return the serialized connector config payload for a specific stage."""
+    if transfer_config is None:
+        return {}
+
+    try:
+        return get_connectors_config_for_stage(transfer_config, stage_id)
+    except Exception as exc:  # pragma: no cover - defensive logging
+        logger.warning(
+            "Failed to build connector config for stage %s: %s. Using IPC fallback.",
+            stage_id,
+            exc,
+        )
+        return {}
+
+
+def build_stage_connectors(
+    stage_id: int,
+    connectors_config: dict[str, Any],
+) -> Optional[dict[tuple[str, str], Any]]:
+    """Instantiate OmniConnectors for a stage based on config."""
+    if not connectors_config:
+        return {}
+
+    logger.info(
+        "[Stage-%s] Initializing OmniConnectors with config keys: %s",
+        stage_id,
+        list(connectors_config.keys()),
+    )
+
+    from .factory import OmniConnectorFactory
+    from .config import ConnectorSpec
+
+    connectors: dict[tuple[str, str], Any] = {}
+    try:
+        # 将字典格式的配置转换为ConnectorSpec对象
+        stage_connector_specs = {}
+        for input_key, config in connectors_config.items():
+            if not input_key.startswith("from_stage_"):
+                continue
+
+            from_stage = input_key.replace("from_stage_", "")
+            spec_dict = config.get("spec", {})
+            if not spec_dict:
+                continue
+
+            connector_spec = ConnectorSpec(
+                name=spec_dict.get("name", "InMemoryConnector"),
+                extra=spec_dict.get("extra", {}),
+            )
+            stage_connector_specs[(str(from_stage), str(stage_id))] = connector_spec
+
+        # 使用统一的连接器创建逻辑
+        connectors = create_connectors_from_config(stage_connector_specs)
+    except Exception as exc:  # pragma: no cover - defensive logging
+        logger.exception("[Stage-%s] Failed to initialize connectors: %s", stage_id, exc)
+        return None
+
+    return connectors
@@ -90,7 +90,7 @@ class AsyncOmni(EngineClient):
        stage_configs_path: Optional[str] = None,
        log_stats: bool = False,
        log_file: Optional[str] = None,
-        init_sleep_seconds: int = 30,
+        init_sleep_seconds: int = 20,
        shm_threshold_bytes: int = 65536,
        batch_timeout: int = 10,
        init_timeout: int = 60000,
@@ -1,38 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import os

 from vllm_omni.diffusion.omni_diffusion import OmniDiffusion
 from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
 from vllm_omni.entrypoints.omni_llm import OmniLLM


-def _dummy_snapshot_download(model_id):
-    return model_id
-
-
-def omni_snapshot_download(model_id) -> str:
-    # 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)
-    else:
-        return _dummy_snapshot_download(model_id)
-
-
 class Omni:
    """Unified entrypoint for both LLM and Diffusion models for better usability."""

    def __init__(self, *args, **kwargs):
        model = args[0] if args else kwargs.get("model", "")
-        assert model != "", "Null model id detected, please specify a model id."
-        model = omni_snapshot_download(model)
-        if args:
-            args[0] = model
-        elif kwargs.get("model", "") != "":
-            kwargs["model"] = model
        if is_diffusion_model(model):
            self.instance: OmniLLM | OmniDiffusion = OmniDiffusion(*args, **kwargs)
        else:
@@ -1,6 +1,8 @@
 import multiprocessing as mp
 import os
+import sys
 import time
+import uuid
 from collections.abc import Sequence
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Any, Optional, Union
@@ -20,6 +22,11 @@ from vllm.utils import Counter
 from vllm.v1.engine.llm_engine import LLMEngine

 # Internal imports (our code)
+from vllm_omni.distributed.connectors.adapter import try_send_via_connector
+from vllm_omni.distributed.connectors import (
+    initialize_orchestrator_connectors,
+    get_stage_connector_config,
+)
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.engine.processor import OmniProcessor
@@ -31,9 +38,14 @@ from vllm_omni.entrypoints.log_utils import (
 )
 from vllm_omni.entrypoints.omni_stage import OmniStage
 from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode
+from vllm_omni.entrypoints.stage_utils import _to_dict
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.stage_utils import serialize_obj as _set
-from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml
+from vllm_omni.entrypoints.utils import (
+    load_stage_configs_from_model,
+    load_stage_configs_from_yaml,
+    resolve_model_config_path,
+)
 from vllm_omni.outputs import OmniRequestOutput

 logger = init_logger(__name__)
@@ -84,14 +96,23 @@ class OmniLLM:
        init_timeout: int = 300,
        **kwargs: Any,
    ):
+        self.worker_backend = kwargs.get("worker_backend", "process")
+        self.ray_address = kwargs.get("ray_address", None)
        self.batch_timeout = batch_timeout
        self._enable_stats: bool = bool(log_stats)
+        self._ray_pg = None
+
        # Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent.
        if stage_configs_path is None:
-            self.stage_configs = load_stage_configs_from_model(model)
+            self.config_path = resolve_model_config_path(model)
+            self.stage_configs = load_stage_configs_from_yaml(self.config_path)
        else:
+            self.config_path = stage_configs_path
            self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)

+        # Initialize connectors
+        self.omni_transfer_config,self.connectors = initialize_orchestrator_connectors(self.config_path)
+
        # Optional file handler for orchestrator
        self._log_file = log_file

@@ -124,36 +145,140 @@ class OmniLLM:
        self.stage_list = [st for _, st in results]
        logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))

-        self._ctx = mp.get_context("spawn")
+        if self.worker_backend == "ray":
+            try:
+                import ray
+                from ray.util.queue import Queue as RayQueue
+            except ImportError:
+                raise ImportError("ray is required for worker_backend='ray'")
+            self._queue_cls = lambda: RayQueue(maxsize=0)
+        else:
+            self._ctx = mp.get_context("spawn")
+            self._queue_cls = lambda: self._ctx.Queue(maxsize=0)
+
        self._stage_in_queues: list[mp.Queue] = []
        self._stage_out_queues: list[mp.Queue] = []
        self._init_sleep_seconds = max(0, int(init_sleep_seconds))
        self._shm_threshold_bytes = max(0, int(shm_threshold_bytes))
-        self._start_stage_processes(model)
+        self._start_stages(model)
        # Wait for all stages to report readiness before seeding
        self._stages_ready: set[int] = set()
        self._wait_for_stages_ready(timeout=init_timeout)

-    def _start_stage_processes(self, model: str) -> None:
+    def _start_stages(self, model: str) -> None:
+        if self.worker_backend == "ray":
+            self._start_stage_actors_ray(model)
+        else:
+            self._start_stage_processes_mp(model)
+
+    def _start_stage_processes_mp(self, model: str) -> None:
        for stage_id, stage in enumerate(self.stage_list):
            # Use unbounded queues to avoid deadlock when seeding many requests
-            in_q: mp.Queue = self._ctx.Queue(maxsize=0)
-            out_q: mp.Queue = self._ctx.Queue(maxsize=0)
+            in_q: mp.Queue = self._queue_cls()
+            out_q: mp.Queue = self._queue_cls()
            self._stage_in_queues.append(in_q)
            self._stage_out_queues.append(out_q)

            # Attach queues and start Stage-owned worker process
            stage.attach_queues(in_q, out_q)
+
+            # Build connectors config for this stage
+            stage_connectors_config = get_stage_connector_config(
+                self.omni_transfer_config,
+                stage_id,
+            )
+
            stage.init_stage_worker(
                model,
                log_file=self._log_file,
                shm_threshold_bytes=self._shm_threshold_bytes,
                ctx=self._ctx,
                batch_timeout=self.batch_timeout,
+                connectors_config=stage_connectors_config,
            )
            logger.debug("[Orchestrator] Stage-%s process started", stage_id)
            time.sleep(self._init_sleep_seconds)

+    def _start_stage_actors_ray(self, model: str) -> None:
+        import ray
+        import os
+        from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+        # Initialize Ray Cluster
+        if not ray.is_initialized():
+             # Create a dummy parallel config just to reuse vLLM's init logic if needed
+             # or just call ray.init()
+             # Pass current PYTHONPATH to workers to ensure they can find vllm_omni
+             runtime_env = {
+                 "env_vars": {"PYTHONPATH": os.environ.get("PYTHONPATH", "")}
+             }
+             ray.init(address=self.ray_address, ignore_reinit_error=True, runtime_env=runtime_env)
+
+        # Create Placement Group
+        # We allocate 1 bundle per stage. Assume 1 GPU per stage for now.
+        # Must include CPU=1 because default Actor request includes 1 CPU.
+        bundles = [{"GPU": 1.0, "CPU": 1.0} for _ in self.stage_list]
+        self._ray_pg = ray.util.placement_group(bundles, strategy="SPREAD")
+        ray.get(self._ray_pg.ready())
+        logger.info("[Orchestrator] Ray Placement Group created")
+
+        # Define Actor Wrapper
+        @ray.remote(num_gpus=1)
+        class OmniStageRayWorker:
+            def run(self, func, *args, **kwargs):
+                return func(*args, **kwargs)
+
+        from vllm_omni.entrypoints.omni_stage import _stage_worker
+
+        for stage_id, stage in enumerate(self.stage_list):
+            in_q = self._queue_cls()
+            out_q = self._queue_cls()
+            self._stage_in_queues.append(in_q)
+            self._stage_out_queues.append(out_q)
+            stage.attach_queues(in_q, out_q)
+
+            stage_connectors_config = get_stage_connector_config(
+                self.omni_transfer_config,
+                stage_id,
+            )
+
+            # Prepare payload
+            # Reuse logic from OmniStage.init_stage_worker but do not start process
+            engine_args = _to_dict(stage.engine_args)
+            runtime_cfg = _to_dict(getattr(stage.stage_config, "runtime", {}))
+            stage_payload = {
+                "stage_id": stage.stage_id,
+                "engine_args": engine_args,
+                "runtime": runtime_cfg,
+                "shm_threshold_bytes": sys.maxsize, # Disable SHM for Ray mode by setting huge threshold
+                "connectors_config": stage_connectors_config or {},
+            }
+
+            # Start Actor
+            # Inject PYTHONPATH again for the specific actor to be safe
+            worker_actor = OmniStageRayWorker.options(
+                scheduling_strategy=PlacementGroupSchedulingStrategy(
+                    placement_group=self._ray_pg,
+                    placement_group_bundle_index=stage_id
+                ),
+                runtime_env={"env_vars": {"PYTHONPATH": os.environ.get("PYTHONPATH", "")}}
+            ).remote()
+
+            # Start worker loop asynchronously
+            worker_actor.run.remote(
+                _stage_worker,
+                model=model,
+                stage_payload=stage_payload,
+                in_q=in_q,
+                out_q=out_q,
+                log_file=self._log_file,
+                batch_timeout=self.batch_timeout,
+            )
+            
+            stage._ray_actor = worker_actor
+            logger.debug("[Orchestrator] Stage-%s Ray Actor started", stage_id)
+            time.sleep(self._init_sleep_seconds)
+
    def close(self) -> None:
        """Close all stage processes and clean up resources.

@@ -171,10 +296,21 @@ class OmniLLM:
                )
        for stage in self.stage_list:
            try:
-                stage.stop_stage_worker()
+                if self.worker_backend == "ray" and hasattr(stage, "_ray_actor"):
+                    import ray
+                    # Ray kill
+                    ray.kill(stage._ray_actor)
+                else:
+                    stage.stop_stage_worker()
            except Exception as e:
                logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)

+        if self.worker_backend == "ray" and self._ray_pg:
+            try:
+                ray.util.remove_placement_group(self._ray_pg)
+            except Exception as e:
+                logger.warning("[Orchestrator] Failed to remove placement group: %s", e)
+
    def __del__(self) -> None:  # best-effort
        try:
            self.close()
@@ -237,11 +373,14 @@ class OmniLLM:
        # Orchestrator keeps stage objects for input derivation
        num_stages = len(self.stage_list)

-        # Map from request_id to original prompt
-        request_id_to_prompt: dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)}
+        # Generate globally unique request IDs and map them to original prompts
+        request_ids: list[str] = [f"{i}_{uuid.uuid4()}" for i in range(len(request_prompts))]
+        request_id_to_prompt: dict[str, PromptType] = {
+            rid: p for rid, p in zip(request_ids, request_prompts)
+        }

        # Track per-request start time for end-to-end timing
-        _req_start_ts: dict[int, float] = {}
+        _req_start_ts: dict[str, float] = {}
        _wall_start_ts: float = time.time()

        # Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage)
@@ -387,51 +526,90 @@ class OmniLLM:
                        )
                        continue
                    sp_next: SamplingParams = sampling_params_list[next_stage_id]  # type: ignore[index]
-                    try:
-                        # Measure transfer size and time (encode + enqueue)
-                        size_bytes = 0
+
+                    # Check if we have a connector for this edge
+                    connector_key = (str(stage_id), str(next_stage_id))
+                    connector = self.connectors.get(connector_key)
+                    print(f"connectors: {self.connectors}")
+
+                    sent_via_connector = False
+                    if connector:
+                        sent_via_connector = try_send_via_connector(
+                            connector=connector,
+                            stage_id=stage_id,
+                            next_stage_id=next_stage_id,
+                            req_id=req_id,
+                            next_inputs=next_inputs,
+                            sampling_params=sp_next,
+                            original_prompt=request_id_to_prompt[req_id],
+                            next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit,
+                            metrics=metrics
+                        )
+
+                    if not sent_via_connector:
+                        # Use original queue mechanism
                        try:
-                            size_bytes = len(_set(next_inputs))
-                        except Exception:
-                            size_bytes = 0
-                        t0 = time.time()
-                        ipc_payload = _encode(
-                            next_inputs,
-                            getattr(self, "_shm_threshold_bytes", 65536),
-                            obj_key="engine_inputs",
-                            shm_key="engine_inputs_shm",
-                        )
-                        ipc_payload.update(
-                            {
-                                "request_id": req_id,
-                                "sampling_params": sp_next,
-                                "sent_ts": time.time(),
-                            }
-                        )
-                        self.stage_list[next_stage_id].submit(ipc_payload)
-                        t1 = time.time()
-                        tx_ms = (t1 - t0) * 1000.0
-                        metrics.on_forward(
-                            stage_id,
-                            next_stage_id,
-                            req_id,
-                            int(size_bytes),
-                            float(tx_ms),
-                            bool("engine_inputs_shm" in ipc_payload),
-                        )
-                    except Exception as e:
-                        logger.warning(
-                            "[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload",
-                            req_id,
-                            e,
-                        )
-                        self.stage_list[next_stage_id].submit(
-                            {
-                                "request_id": req_id,
-                                "engine_inputs": next_inputs,
-                                "sampling_params": sp_next,
-                            }
-                        )
+                            if self.worker_backend == "ray":
+                                # Ray mode: Skip SHM fallback, just send object via Queue (Ray handles serialization)
+                                # This prevents SHM access errors across nodes
+                                ipc_payload = {
+                                    "request_id": req_id,
+                                    "engine_inputs": next_inputs,
+                                    "sampling_params": sp_next,
+                                    "sent_ts": time.time(),
+                                }
+                                size_bytes = 0 # Approximation or calculate size if needed
+                                try:
+                                    # Approximate size for stats
+                                    size_bytes = len(_set(next_inputs))
+                                except:
+                                    pass
+                                t0 = time.time()
+                                # Direct submit to Ray Queue
+                            else:
+                                size_bytes = 0
+                                try:
+                                    size_bytes = len(_set(next_inputs))
+                                except Exception:
+                                    size_bytes = 0
+                                t0 = time.time()
+                                ipc_payload = _encode(
+                                    next_inputs,
+                                    getattr(self, "_shm_threshold_bytes", 65536),
+                                    obj_key="engine_inputs",
+                                    shm_key="engine_inputs_shm",
+                                )
+                                ipc_payload.update(
+                                    {
+                                        "request_id": req_id,
+                                        "sampling_params": sp_next,
+                                        "sent_ts": time.time(),
+                                    }
+                                )
+                            self.stage_list[next_stage_id].submit(ipc_payload)
+                            t1 = time.time()
+                            tx_ms = (t1 - t0) * 1000.0
+                            metrics.on_forward(
+                                stage_id,
+                                next_stage_id,
+                                req_id,
+                                int(size_bytes),
+                                float(tx_ms),
+                                bool("engine_inputs_shm" in ipc_payload),
+                            )
+                        except Exception as e:
+                            logger.warning(
+                                "[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload",
+                                req_id,
+                                e,
+                            )
+                            self.stage_list[next_stage_id].submit(
+                                {
+                                    "request_id": req_id,
+                                    "engine_inputs": next_inputs,
+                                    "sampling_params": sp_next,
+                                }
+                            )
                    logger.debug(
                        "[Orchestrator] Forwarded request %s to stage-%s",
                        req_id,
@@ -163,6 +163,7 @@ class OmniStage:
        shm_threshold_bytes: int = 65536,
        ctx: Optional[mp.context.BaseContext] = None,
        batch_timeout: int = 10,
+        connectors_config: Optional[dict] = None,  # New parameter
    ) -> None:
        """Initialize and start the stage worker process.

@@ -192,6 +193,7 @@ class OmniStage:
            "engine_args": engine_args,
            "runtime": runtime_cfg,
            "shm_threshold_bytes": self._shm_threshold_bytes,
+            "connectors_config": connectors_config or {},
        }
        if is_async:
            self._proc = ctx.Process(
@@ -340,6 +342,8 @@ def _stage_worker(
        log_stage_batch_stats,
        log_stage_running_avg,
    )
+    from vllm_omni.distributed.connectors.adapter import try_recv_via_connector
+    from vllm_omni.distributed.connectors.utils import build_stage_connectors
    from vllm_omni.entrypoints.omni_llm import OmniStageLLM
    # no inline JSONL/serialization imports; logging handled by utilities

@@ -347,6 +351,7 @@ def _stage_worker(
    engine_args = stage_payload.get("engine_args", {})
    runtime_cfg = stage_payload.get("runtime", {})
    shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
+    connectors_config = stage_payload.get("connectors_config", {})

    # Per-stage logger: clear inherited handlers to avoid broken parent streams
    try:
@@ -440,6 +445,18 @@ def _stage_worker(
    )
    stage_engine = OmniStageLLM(model=model, **engine_args)
    _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
+
+    # Initialize OmniConnectors if configured
+    connectors = {}
+    if connectors_config:
+        built_connectors = build_stage_connectors(
+            stage_id=stage_id,
+            connectors_config=connectors_config,
+        )
+        if built_connectors is None:
+            return
+        connectors = built_connectors
+
    # Signal readiness to orchestrator
    try:
        out_q.put({"type": "stage_ready", "stage_id": stage_id})
@@ -496,9 +513,21 @@ def _stage_worker(
                    _in_flight_ms_by_rid[rid] = 0.0
            except Exception:
                _in_flight_ms_by_rid[rid] = 0.0
-            ein, _rx_metrics = maybe_load_from_ipc_with_metrics(t, obj_key="engine_inputs", shm_key="engine_inputs_shm")
-            _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
-            _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
+
+            # Resolve input data (via Connector or IPC)
+            ein, _rx_metrics = try_recv_via_connector(
+                task=t,
+                connectors=connectors,
+                stage_id=stage_id
+            )
+
+            if ein is None or _rx_metrics is None:
+                ein, _rx_metrics = maybe_load_from_ipc_with_metrics(task, obj_key="engine_inputs", shm_key="engine_inputs_shm")
+
+            if _rx_metrics:
+                _rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
+                _rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
+
            batch_request_ids.append(rid)
            if isinstance(ein, list):
                batch_engine_inputs.extend(ein)
@@ -10,6 +10,41 @@ from vllm_omni.utils import detect_device_type
 PROJECT_ROOT = Path(__file__).parent.parent.parent


+def resolve_model_config_path(model: str) -> str:
+    """Resolve the stage config file path from the model name.
+
+    Resolves stage configuration path based on the model type and device type.
+    First tries to find a device-specific YAML file from stage_configs/{device_type}/
+    directory. If not found, falls back to the default config file.
+
+    Args:
+        model: Model name or path (used to determine model_type)
+
+    Returns:
+        String path to the stage configuration file
+
+    Raises:
+        FileNotFoundError: If no stage config file exists for the model type
+    """
+    hf_config = get_config(model, trust_remote_code=True)
+    model_type = hf_config.model_type
+    device_type = detect_device_type()
+
+    # Try device-specific config first
+    if device_type != "cuda":
+        device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml"
+        device_config_path = PROJECT_ROOT / device_config_file
+        if os.path.exists(device_config_path):
+            return str(device_config_path)
+
+    # Fall back to default config
+    stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
+    stage_config_path = PROJECT_ROOT / stage_config_file
+    if not os.path.exists(stage_config_path):
+        raise FileNotFoundError(f"Stage config file {stage_config_path} not found")
+    return str(stage_config_path)
+
+
 def load_stage_configs_from_model(model: str) -> list:
    """Load stage configurations from model's default config file.

@@ -26,24 +61,8 @@ def load_stage_configs_from_model(model: str) -> list:
    Raises:
        FileNotFoundError: If no stage config file exists for the model type
    """
-    hf_config = get_config(model, trust_remote_code=True)
-    model_type = hf_config.model_type
-    device_type = detect_device_type()
-
-    # Try device-specific config first
-    if device_type != "cuda":
-        device_config_file = f"vllm_omni/model_executor/stage_configs/{device_type}/{model_type}.yaml"
-        device_config_path = PROJECT_ROOT / device_config_file
-        if os.path.exists(device_config_path):
-            stage_configs = load_stage_configs_from_yaml(config_path=str(device_config_path))
-            return stage_configs
-
-    # Fall back to default config
-    stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
-    stage_config_path = PROJECT_ROOT / stage_config_file
-    if not os.path.exists(stage_config_path):
-        raise FileNotFoundError(f"Stage config file {stage_config_path} not found")
-    stage_configs = load_stage_configs_from_yaml(config_path=str(stage_config_path))
+    stage_config_path = resolve_model_config_path(model)
+    stage_configs = load_stage_configs_from_yaml(config_path=stage_config_path)
    return stage_configs


@@ -1,6 +1,4 @@
 # stage config for running qwen2.5-omni with architecture of OmniLLM.
-
-# The following config has been verified on 1x H100-80G GPU.
 stage_args:
  - stage_id: 0
    runtime:
@@ -28,6 +26,9 @@ stage_args:
      seed: 42
      detokenize: True
      repetition_penalty: 1.1
+    # Distributed connector configuration (optional)
+    output_connectors:
+      to_stage_1: connector_thinker2talker # thinker -> talker: use local mooncake
  - stage_id: 1
    runtime:
      process: true
@@ -54,17 +55,22 @@ stage_args:
      detokenize: True
      repetition_penalty: 1.05
      stop_token_ids: [8294]
+    # Distributed connector configuration (optional)
+    input_connectors:
+      from_stage_0: connector_thinker2talker # talker <- thinker: use local mooncake
+    output_connectors:
+      to_stage_2: connector_talker2generator # talker -> code2wav: use remote mooncake (example)
  - stage_id: 2
    runtime:
      process: true
-      devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
+      devices: "2"            # Example: use a different GPU than the previous stage; use "0" if single GPU
      max_batch_size: 1
    engine_args:
      model_stage: code2wav
      model_arch: Qwen2_5OmniForConditionalGeneration
      worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
      scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
-      gpu_memory_utilization: 0.15
+      gpu_memory_utilization: 0.3
      enforce_eager: true
      trust_remote_code: true
      enable_prefix_caching: false
@@ -80,6 +86,9 @@ stage_args:
      seed: 42
      detokenize: True
      repetition_penalty: 1.1
+    # Distributed connector configuration (optional)
+    input_connectors:
+      from_stage_1: connector_talker2generator # code2wav <- talker: use remote mooncake

 # Top-level runtime config (concise): default windows and stage edges
 runtime:
@@ -87,6 +96,31 @@ runtime:
  defaults:
    window_size: -1             # Simplified: trigger downstream only after full upstream completion
    max_inflight: 1             # Simplified: process serially within each stage
+
+  # Distributed connectors configuration (optional)
+  connectors:
+    # Local mooncake for same-node communication
+    connector_thinker2talker:
+      name: MooncakeConnector
+      extra:
+        host: "127.0.0.1"
+        metadata_server: "http://10.90.67.86:8080/metadata"
+        master: "10.90.67.86:50051"
+        segment: 512000000    # 512MB
+        localbuf: 64000000     # 64MB
+        proto: "tcp"
+
+    # Remote mooncake for cross-node communication (example)
+    connector_talker2generator:
+      name: MooncakeConnector
+      extra:
+        host: "127.0.0.1"
+        metadata_server: "http://10.90.67.86:8080/metadata"
+        master: "10.90.67.86:50051"
+        segment: 1073741824    # 1GB
+        localbuf: 134217728     # 128MB
+        proto: "tcp"
+
  edges:
    - from: 0                   # thinker → talker: trigger only after receiving full input (-1)
      to: 1
@@ -3,25 +3,24 @@
 # Stage 1: Talker (text embeddings → 8-layer RVQ codec codes)
 # Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)

-# The following config has been verified on 2x H100-80G GPUs.
 stage_args:
  - stage_id: 0
    runtime:
-      devices: "0,1"
+      devices: "0"
      max_batch_size: 1
    engine_args:
      model_stage: thinker
      model_arch: Qwen3OmniMoeForConditionalGeneration
      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
-      gpu_memory_utilization: 0.6
+      gpu_memory_utilization: 0.8
      enforce_eager: true
      trust_remote_code: true
      engine_output_type: latent  # Output hidden states for talker
      distributed_executor_backend: "mp"
      enable_prefix_caching: false
      hf_config_name: thinker_config
-      tensor_parallel_size: 2
+      # tensor_parallel_size: 2  # Commented out - using single GPU for thinker
    final_output: true
    final_output_type: text
    is_comprehension: true
@@ -43,7 +42,7 @@ stage_args:
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
-       gpu_memory_utilization: 0.3
+       gpu_memory_utilization: 0.5
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: latent  # Output codec codes for code2wav
@@ -66,7 +65,7 @@ stage_args:

  - stage_id: 2
    runtime:
-      devices: "0"
+      devices: "1"
      max_batch_size: 1
    engine_args:
      model_stage: code2wav
@@ -15,9 +15,6 @@ from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
 from vllm_omni.request import OmniRequest

 for module_name, module in sys.modules.items():
-    # only do patch on module of vllm, pass others
-    if "vllm" not in module_name:
-        continue
    if hasattr(module, "EngineCoreOutput") and module.EngineCoreOutput == _OriginalEngineCoreOutput:
        module.EngineCoreOutput = OmniEngineCoreOutput
    if hasattr(module, "EngineCoreOutputs") and module.EngineCoreOutputs == _OriginalEngineCoreOutputs:
Author	SHA1	Message	Date
wzliu	26cf9ef4ed	update	2025-12-01 10:49:39 +08:00
wzliu	33e03b44c5	added omni mooncake connectors to distributed directory, integrate omni connector to omni-llm, rebased with main (twice). ray + mooncake connector work on 2 nodes, but needes several updates\n1.start ray inside vllm-omni \n2.placement group optimize \n3.fix close	2025-12-01 10:45:26 +08:00