[BugFix]: fix a lot of bug (#1565)

Signed-off-by: princepride <wangzhipeng628@gmail.com>

.buildkite/pipeline.yml

@@ -3,7 +3,7 @@ steps:
     key: image-build
     commands:
       - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-      - "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
+      - "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
       - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
       - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
     agents:

.buildkite/test-amd.yaml

@@ -33,6 +33,7 @@ steps:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+    - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
 
 - label: "Diffusion Cache Backend Test"
   timeout_in_minutes: 15
@@ -88,6 +89,7 @@ steps:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+    - pytest -s -v tests/engine/test_async_omni_engine_abort.py
 
 - label: "Omni Model Test Qwen3-Omni"
   timeout_in_minutes: 15
@@ -101,7 +103,7 @@ steps:
     - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
     - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
     - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
-    - pytest -s -v tests/engine/test_async_omni_engine_abort.py
+
 
 - label: "Diffusion Image Edit Test"
   timeout_in_minutes: 15

.buildkite/test-merge.yml

@@ -24,7 +24,7 @@ steps:
       queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
     plugins:
       - docker#v5.2.0:
-          image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
           always-pull: true
           propagate-environment: true
           environment:
@@ -155,41 +155,31 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Benchmark & Engine Test with H100"
-    timeout_in_minutes: 15
+  - label: "Benchmark & Engine Test"
     depends_on: upload-merge-pipeline
     commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/benchmarks/test_serve_cli.py
-      - pytest -s -v tests/engine/test_async_omni_engine_abort.py
+      - |
+        timeout 15m bash -c '
+                export VLLM_WORKER_MULTIPROC_METHOD=spawn
+                set +e
+                pytest -s -v tests/benchmarks/test_serve_cli.py
+                EXIT1=$$?
+                pytest -s -v tests/engine/test_async_omni_engine_abort.py
+                EXIT2=$$?
+                exit $$((EXIT1 | EXIT2))
+        '
     agents:
-      queue: "mithril-h100-pool"
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 2
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Omni Model Test"
     timeout_in_minutes: 15

.buildkite/test-nightly.yml

@@ -5,15 +5,20 @@ steps:
     if: build.env("NIGHTLY") == "1"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py -m "advanced_model" --run-level "advanced_model"
-      - pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
+      - |
+        set +e
+        pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py -m "advanced_model" --run-level "advanced_model"
+        EXIT1=$$?
+        pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
+        EXIT2=$$?
+        exit $$((EXIT1 | EXIT2))
     agents:
       queue: "mithril-h100-pool"
     plugins:
       - kubernetes:
           podSpec:
             containers:
-              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
                 resources:
                   limits:
                     nvidia.com/gpu: 2
@@ -69,7 +74,7 @@ steps:
       - kubernetes:
           podSpec:
             containers:
-              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
                 resources:
                   limits:
                     nvidia.com/gpu: 2

.buildkite/test-ready.yml

@@ -2,7 +2,7 @@ steps:
   - label: "Simple Unit Test"
     depends_on: upload-ready-pipeline
     commands:
-      - "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+      - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
     agents:
       queue: "gpu_1_queue"
     plugins:
@@ -16,12 +16,11 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Model Test"
-    timeout_in_minutes: 20
     depends_on: upload-ready-pipeline
     commands:
-      - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
     agents:
-      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+      queue: "gpu_1_queue"
     plugins:
       - docker#v5.2.0:
           image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -33,11 +32,17 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Model CPU offloading Test"
-    timeout_in_minutes: 20
     depends_on: upload-ready-pipeline
     commands:
-      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
-      - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+      - |
+        timeout 20m bash -c '
+          set +e
+          pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+          EXIT1=$$?
+          pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+          EXIT2=$$?
+          exit $$((EXIT1 | EXIT2))
+        '
     agents:
       queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
     plugins:
@@ -51,10 +56,9 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Audio Generation Model Test"
-    timeout_in_minutes: 20
     depends_on: upload-ready-pipeline
     commands:
-      - pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
     agents:
       queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
     plugins:
@@ -68,10 +72,9 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Cache Backend Test"
-    timeout_in_minutes: 15
     depends_on: upload-ready-pipeline
     commands:
-      - pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
+      - timeout 15m pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
     agents:
       queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
     plugins:
@@ -85,10 +88,9 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Sequence Parallelism Test"
-    timeout_in_minutes: 20
     depends_on: upload-ready-pipeline
     commands:
-      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:
@@ -103,10 +105,9 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion GPU Worker Test"
-    timeout_in_minutes: 20
     depends_on: upload-ready-pipeline
     commands:
-      - pytest -s -v tests/diffusion/test_diffusion_worker.py
+      - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:
@@ -120,49 +121,43 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  # - label: "Benchmark & Engine Test with H100"
-  #   timeout_in_minutes: 15
-  #   depends_on: upload-ready-pipeline
-  #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - pytest -s -v tests/benchmarks/test_serve_cli.py
-  #     - pytest -s -v tests/engine/test_async_omni_engine_abort.py
-  #   agents:
-  #     queue: "mithril-h100-pool"
-  #   plugins:
-  #     - kubernetes:
-  #         podSpec:
-  #           containers:
-  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-  #               resources:
-  #                 limits:
-  #                   nvidia.com/gpu: 2
-  #               volumeMounts:
-  #                 - name: devshm
-  #                   mountPath: /dev/shm
-  #                 - name: hf-cache
-  #                   mountPath: /root/.cache/huggingface
-  #               env:
-  #                 - name: HF_HOME
-  #                   value: /root/.cache/huggingface
-  #           nodeSelector:
-  #             node.kubernetes.io/instance-type: gpu-h100-sxm
-  #           volumes:
-  #             - name: devshm
-  #               emptyDir:
-  #                 medium: Memory
-  #             - name: hf-cache
-  #               hostPath:
-  #                 path: /mnt/hf-cache
-  #                 type: DirectoryOrCreate
 
-  - label: "Omni Model Test"
-    timeout_in_minutes: 17
+  - label: "Benchmark & Engine Test"
     depends_on: upload-ready-pipeline
     commands:
-      - export VLLM_LOGGING_LEVEL=DEBUG
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+      - |
+        timeout 15m bash -c '
+                export VLLM_WORKER_MULTIPROC_METHOD=spawn
+                set +e
+                pytest -s -v tests/benchmarks/test_serve_cli.py
+                EXIT1=$$?
+                pytest -s -v tests/engine/test_async_omni_engine_abort.py
+                EXIT2=$$?
+                exit $$((EXIT1 | EXIT2))
+        '
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+
+  - label: "Omni Model Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 17m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+        '
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:
@@ -176,20 +171,22 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   # - label: "Omni Model Test with H100"
-  #   timeout_in_minutes: 20
   #   depends_on: upload-ready-pipeline
   #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-  #     # - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-  #     - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
+  #     - |
+  #       timeout 20m bash -c '
+  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #         export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+  #         # - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+  #         pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
+  #       '
   #   agents:
   #     queue: "mithril-h100-pool"
   #   plugins:
   #     - kubernetes:
   #         podSpec:
   #           containers:
-  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
   #               resources:
   #                 limits:
   #                   nvidia.com/gpu: 2
@@ -213,12 +210,14 @@ steps:
   #                 type: DirectoryOrCreate
 
   - label: "Qwen3-TTS E2E Test"
-    timeout_in_minutes: 20
     depends_on: upload-ready-pipeline
     commands:
-      - export VLLM_LOGGING_LEVEL=DEBUG
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
+      - |
+        timeout 20m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
+        '
     agents:
       queue: "gpu_4_queue"
     plugins:
@@ -233,18 +232,20 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   # - label: "Diffusion Image Edit Test with H100 (1 GPU)"
-  #   timeout_in_minutes: 20
   #   depends_on: upload-ready-pipeline
   #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+  #     - |
+  #       timeout 20m bash -c '
+  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #         pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+  #       '
   #   agents:
   #     queue: "mithril-h100-pool"
   #   plugins:
   #     - kubernetes:
   #         podSpec:
   #           containers:
-  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
   #               resources:
   #                 limits:
   #                   nvidia.com/gpu: 1
@@ -268,18 +269,20 @@ steps:
   #                 type: DirectoryOrCreate
 
   # - label: "Bagel Text2Img Model Test with H100"
-  #   timeout_in_minutes: 30
   #   depends_on: upload-ready-pipeline
   #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
+  #     - |
+  #       timeout 30m bash -c '
+  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #         pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
+  #       '
   #   agents:
   #     queue: "mithril-h100-pool"
   #   plugins:
   #     - kubernetes:
   #         podSpec:
   #           containers:
-  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
   #               resources:
   #                 limits:
   #                   nvidia.com/gpu: 1

.github/PULL_REQUEST_TEMPLATE.md

@@ -12,10 +12,10 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
 <summary> Essential Elements of an Effective PR Description Checklist </summary>
 
 - [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
-- [ ] The test plan. Please providing the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
-- [ ] The test results. Please pasting the results comparison before and after, or e2e results.
+- [ ] The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
+- [ ] The test results. Please paste the results comparison before and after, or the e2e results.
 - [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. **Please run `mkdocs serve` to sync the documentation editions to `./docs`.**
-- [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft.
+- [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
 </details>
 
 **BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)

.github/workflows/comfyui-publish.yml

@@ -1,27 +0,0 @@
-name: ComfyUI Integration -- Publish to Comfy registry
-on:
-  workflow_dispatch:
-  push:
-    tags:
-      - "*"
-    paths:
-      - "apps/ComfyUI-vLLM-Omni/**"
-
-permissions:
-  issues: write
-
-defaults:
-  run:
-    working-directory: apps/ComfyUI-vLLM-Omni
-
-jobs:
-  publish-node:
-    name: Publish Custom Node to registry
-    runs-on: ubuntu-latest
-    steps:
-      - name: ♻️ Check out code
-        uses: actions/checkout@v4
-      - name: 📦 Publish Custom Node
-        uses: Comfy-Org/publish-node-action@main
-        with:
-          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.github/workflows/comfyui-validate.yml

@@ -1,40 +0,0 @@
-name: ComfyUI Integration -- Validate backwards compatibility
-
-on:
-  pull_request:
-    branches:
-      - master
-      - main
-    paths:
-      - "apps/ComfyUI-vLLM-Omni/**"
-
-defaults:
-  run:
-    working-directory: apps/ComfyUI-vLLM-Omni
-
-jobs:
-  check-base-path:
-    runs-on: ubuntu-latest
-    outputs:
-      exists: ${{ steps.check.outputs.exists }}
-    defaults:
-      run:
-        working-directory: .
-    steps:
-      - uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event.pull_request.base.sha }}
-          fetch-depth: 1
-      - id: check
-        run: |
-          if [ -d "apps/ComfyUI-vLLM-Omni" ]; then
-            echo "exists=true" >> "$GITHUB_OUTPUT"
-          else
-            echo "exists=false" >> "$GITHUB_OUTPUT"
-          fi
-  validate:
-    needs: check-base-path
-    if: needs.check-base-path.outputs.exists == 'true'
-    runs-on: ubuntu-latest
-    steps:
-      - uses: comfy-org/node-diff@main

.gitignore

@@ -243,5 +243,7 @@ Dockerfile.dev
 discussion
 tmp_test
 
+# Auto-generated version file (created by setuptools_scm during build)
+vllm_omni/_version.py
 # output files
 *.wav

apps/ComfyUI-vLLM-Omni/README.md

@@ -47,12 +47,12 @@ If no, check your shell running the ComfyUI process. There may be some error mes
 This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
 
 - **Generate Image** for text-to-image and image-to-image tasks
-- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks
+- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
 - **TTS** and **TTS Voice Clone** for TTS tasks
 
 This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
 
-> [!INFO]
+> [!NOTE]
 > The node UI and feature designs are intended to match vLLM-Omni online serving interfaces. It cannot offer more than what the interfaces support.
 
 To build a simple workflow yourself,
@@ -65,28 +65,16 @@ To build a simple workflow yourself,
     - For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
     - For multi-stage models where all stages are either autoregression or diffusion, you can also connect only a single Sampling Params node, indicating that this set of sampling parameters will be used for all stages.
 
-**The following features are tested**:
-
-- Single-node workflows for
-    - Multimodal Comprehension (e.g., Qwen Omni, BAGEL)
-    - Text-to-Image Generation (e.g., Qwen-Image)
-    - Image-to-Image Generation (e.g., Qwen-Image-Edit)
-    - TTS (e.g., Qwen TTS, including VoiceDesign, VoiceClone, CustomVoice)
-
-**The following features are not currently tested**. They may work or break. You are welcomed to test it out and offer comments.
-
-- Multi-node workflow that connects multiple model services together.
-
 ## Screenshots and Examples
 
-### Multimodal comprehension (e.g., Qwen Omni series, BAGEL)
+### Multimodal understanding (e.g., Qwen Omni series, BAGEL)
 
-(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Annotated Example**)
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Multimodal Understanding**)
 
 <p align="center">
   <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg">
-    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg" width=55%>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-understanding.jpg">
+    <img alt="vLLM-Omni multimodal understanding" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-understanding.jpg" width=55%>
   </picture>
 </p>
 
@@ -98,7 +86,7 @@ You can configure per-stage sampling parameters for multi-stage models.
 <p align="center">
   <picture>
     <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg">
-    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg" width=55%>
+    <img alt="vLLM-Omni multiple stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg" width=55%>
   </picture>
 </p>
 
@@ -109,7 +97,7 @@ You can configure per-stage sampling parameters for multi-stage models.
 <p align="center">
   <picture>
     <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg">
-    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg" width=55%>
+    <img alt="vLLM-Omni image generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg" width=55%>
   </picture>
 </p>
 
@@ -123,13 +111,24 @@ You can configure per-stage sampling parameters for multi-stage models.
 <p align="center">
   <picture>
     <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg">
-    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg" width=55%>
+    <img alt="vLLM-Omni TTS" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg" width=55%>
   </picture>
 </p>
 
 > [!TIP]
 > There is a dedicated node for VoiceClone tasks with reference audio input. Other simple text-to-speech tasks should use the regular TTS node.
 
+### Chaining multiple model services
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Chaining Services**)
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-chaining-services.jpg">
+    <img alt="vLLM-Omni TTS" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-chaining-services.jpg" width=55%>
+  </picture>
+</p>
+
 ## Develop
 
 Follow the [development convention and rules of vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/).

apps/ComfyUI-vLLM-Omni/__init__.py

@@ -6,18 +6,18 @@ __all__ = [
     "WEB_DIRECTORY",
 ]
 
-__author__ = """Zeyu Huang"""
-__email__ = "11222265+fhfuih@users.noreply.github.com"
+__author__ = """vLLM-Omni Team"""
+__email__ = "vllm-omni@vllm.ai"
 __version__ = "0.0.1"
 
 from .comfyui_vllm_omni.nodes import (
     VLLMOmniARSampling,
-    VLLMOmniComprehension,
     VLLMOmniDiffusionSampling,
     VLLMOmniGenerateImage,
     VLLMOmniQwenTTSParams,
     VLLMOmniSamplingParamsList,
     VLLMOmniTTS,
+    VLLMOmniUnderstanding,
     VLLMOmniVoiceClone,
 )
 
@@ -25,7 +25,7 @@ from .comfyui_vllm_omni.nodes import (
 NODE_CLASS_MAPPINGS = {
     # === Generation ===
     "VLLMOmniGenerateImage": VLLMOmniGenerateImage,
-    "VLLMOmniComprehension": VLLMOmniComprehension,
+    "VLLMOmniUnderstanding": VLLMOmniUnderstanding,
     "VLLMOmniTTS": VLLMOmniTTS,
     "VLLMOmniVoiceClone": VLLMOmniVoiceClone,
     # === Params ===
@@ -39,7 +39,7 @@ NODE_CLASS_MAPPINGS = {
 NODE_DISPLAY_NAME_MAPPINGS = {
     # === Generation ===
     "VLLMOmniGenerateImage": "Generate Image",
-    "VLLMOmniComprehension": "Multimodality Comprehension",
+    "VLLMOmniUnderstanding": "Multimodality Understanding",
     "VLLMOmniTTS": "TTS (Text to Speech)",
     "VLLMOmniVoiceClone": "TTS Voice Cloning",
     # === Params ===

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py

@@ -1,4 +1,4 @@
-from typing import Literal, cast
+from typing import Literal
 
 import torch
 from comfy_api.input import AudioInput, VideoInput
@@ -6,7 +6,12 @@ from comfy_api.input import AudioInput, VideoInput
 from .utils.api_client import VLLMOmniClient
 from .utils.logger import get_logger
 from .utils.models import lookup_model_spec
-from .utils.types import AudioFormat
+from .utils.types import (
+    AudioFormat,
+    AutoregressionSamplingParams,
+    DiffusionSamplingParams,
+    QwenTTSModelSpecificParams,
+)
 from .utils.validators import (
     add_sampling_parameters_to_stage,
     validate_model_and_sampling_params_types,
@@ -86,7 +91,10 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
 
         # Prefer DALL-E compatible API for simple (one-stage) diffusion models
         if (spec is None or spec["stages"] == ["diffusion"]) and not is_bagel:
-            sampling_params = cast(dict | None, sampling_params)
+            # The number of sampling parameter groups should have been validated.
+            # Now, simply convert single-item list to dict.
+            if isinstance(sampling_params, list):
+                sampling_params = sampling_params[0]
             if audio is None and image is None and video is None:
                 # No multimodal input --- use DALL-E image generation
                 logger.info("Using DALL-E image generation endpoint")
@@ -133,7 +141,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
         return (output,)
 
 
-class VLLMOmniComprehension(_VLLMOmniGenerateBase):
+class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
     @classmethod
     def INPUT_TYPES(cls):
         return {
@@ -197,7 +205,7 @@ class VLLMOmniComprehension(_VLLMOmniGenerateBase):
             (
                 text_response,
                 _,
-            ) = await client.generate_comprehension_chat_completion(
+            ) = await client.generate_understanding_chat_completion(
                 model=model,
                 prompt=prompt,
                 image=image,
@@ -221,7 +229,7 @@ class VLLMOmniComprehension(_VLLMOmniGenerateBase):
             (
                 text_response,
                 audio,
-            ) = await client.generate_comprehension_chat_completion(
+            ) = await client.generate_understanding_chat_completion(
                 model=model,
                 prompt=prompt,
                 image=image,
@@ -287,15 +295,13 @@ class VLLMOmniTTS(_VLLMOmniGenerateBase):
         logger.info("Got extra kwargs in TTS: %s", kwargs)
 
         is_qwen_tts = "qwen3-tts" in model.lower()
-        extra_params_type = None if model_specific_params is None else model_specific_params["type"]
-        if not is_qwen_tts and extra_params_type == "qwen-tts":
+        if not is_qwen_tts and isinstance(model_specific_params, QwenTTSModelSpecificParams):
             raise ValueError(
                 "You have provided Qwen-specific TTS params."
                 "However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
             )
 
         combined_params = {**kwargs, **(model_specific_params or {})}
-        combined_params.pop("type", None)  # Internal fields in model_specific_params
 
         client = VLLMOmniClient(url)
 
@@ -352,8 +358,7 @@ class VLLMOmniVoiceClone(_VLLMOmniGenerateBase):
         **kwargs,
     ):
         is_qwen_tts = "qwen3-tts" in model.lower()
-        extra_params_type = None if model_specific_params is None else model_specific_params["type"]
-        if not is_qwen_tts and extra_params_type == "qwen-tts":
+        if not is_qwen_tts and isinstance(model_specific_params, QwenTTSModelSpecificParams):
             raise ValueError(
                 "You have provided Qwen-specific TTS params."
                 "However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
@@ -366,7 +371,6 @@ class VLLMOmniVoiceClone(_VLLMOmniGenerateBase):
             **kwargs,
             **(model_specific_params or {}),
         }
-        combined_params.pop("type", None)  # Internal fields in model_specific_params
 
         client = VLLMOmniClient(url)
 
@@ -419,10 +423,7 @@ class VLLMOmniARSampling:
     CATEGORY = "vLLM-Omni/Sampling Params"
 
     def get_params(self, seed, **kwargs):
-        params = {
-            "type": "autoregression",  # for internal use, removed before sending the request
-            **kwargs,
-        }
+        params = AutoregressionSamplingParams(kwargs)
         if seed >= 0:
             params["seed"] = seed
         return (params,)
@@ -479,6 +480,13 @@ class VLLMOmniDiffusionSampling:
                         "tooltip": "Enable VAE slicing for reduced memory usage (slight quality trade-off)",
                     },
                 ),
+                "vae_use_tiling": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": "Enable VAE tiling for reduced memory usage (slight quality trade-off)",
+                    },
+                ),
                 # === Put seed at last. ===
                 # Whenever a field named "seed" is present, ComfyUI adds another field called "control after generate"
                 "seed": (
@@ -499,10 +507,7 @@ class VLLMOmniDiffusionSampling:
     CATEGORY = "vLLM-Omni/Sampling Params"
 
     def get_params(self, seed, **kwargs):
-        params = {
-            "type": "diffusion",  # for internal use, removed before sending the request
-            **kwargs,
-        }
+        params = DiffusionSamplingParams(kwargs)
         if seed >= 0:
             params["seed"] = seed
         return (params,)
@@ -566,4 +571,4 @@ class VLLMOmniQwenTTSParams:
     CATEGORY = "vLLM-Omni/TTS Params"
 
     def get_params(self, **kwargs):
-        return ({"type": "qwen-tts", **kwargs},)
+        return (QwenTTSModelSpecificParams(kwargs),)

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py

@@ -57,16 +57,7 @@ class VLLMOmniClient:
         if negative_prompt:
             payload["negative_prompt"] = negative_prompt
         if sampling_params is not None:
-            # Only select specific sampling params
-            for k in (
-                "n",
-                "num_inference_steps",
-                "guidance_scale",
-                "true_cfg_scale",
-                "vae_use_slicing",
-            ):
-                if k in sampling_params and sampling_params[k] is not None:
-                    payload[k] = sampling_params[k]
+            payload.update(sampling_params)
         logger.debug("img gen payload: %s", payload)
 
         url = self.base_url + "/images/generations"
@@ -138,10 +129,8 @@ class VLLMOmniClient:
         if negative_prompt:
             form.add_field("negative_prompt", negative_prompt)
         if sampling_params is not None:
-            # Only select specific sampling params
-            for k in ("n", "num_inference_steps", "guidance_scale", "true_cfg_scale"):
-                if k in sampling_params and sampling_params[k] is not None:
-                    form.add_field(k, str(sampling_params[k]))
+            for k, v in sampling_params.items():
+                form.add_field(k, str(v))
         if mask is not None:
             mask_filename = "mask.png"
             form.add_field(
@@ -217,7 +206,7 @@ class VLLMOmniClient:
 
         return torch.stack(image_tensors, dim=0)
 
-    async def generate_comprehension_chat_completion(
+    async def generate_understanding_chat_completion(
         self,
         *,
         model: str,
@@ -427,30 +416,28 @@ class VLLMOmniClient:
             message_content.append({"type": "video_url", "video_url": {"url": video_to_base64(video)}})
         messages = [{"role": "user", "content": message_content}]
 
+        payload: dict[str, Any] = {"messages": messages, "model": model}
+        if modalities:
+            payload["modalities"] = modalities
+
         combined_extra_body: dict[str, Any] = {}
         if sampling_params is not None:
             spec, _ = lookup_model_spec(model)
             is_single_sampling_param = isinstance(sampling_params, dict) or len(sampling_params) == 1
 
-            # Exclude internal key
-            if isinstance(sampling_params, dict):
-                sampling_params = {k: v for k, v in sampling_params.items() if k != "type"}
-            else:
-                sampling_params = [{k: v for k, v in sp.items() if k != "type"} for sp in sampling_params]
-
             if (spec is None and is_single_sampling_param) or (spec is not None and spec["stages"] == ["diffusion"]):
                 # Diffusion format: extra_body directly contains sampling params.
-                # Validation should have taken care of matching sampling params' types.
+                # Validation has already taken care of matching sampling params' types and length. Safe to take [0].
                 # * Use this mode if the model is a simple one-stage diffusion model.
                 # * Fallback to this mode if model is not registered and a single sampling param is provided.
                 sampling_params = sampling_params if isinstance(sampling_params, dict) else sampling_params[0]
-                combined_extra_body: dict[str, Any] = {**sampling_params}
+                combined_extra_body: dict[str, Any] = sampling_params.copy()
                 if "n" in combined_extra_body:
-                    combined_extra_body["num_outputs_per_prompt"] = combined_extra_body["n"]
-                    del combined_extra_body["n"]
+                    combined_extra_body["num_outputs_per_prompt"] = combined_extra_body.pop("n")
             else:
-                # Use AR style payload, extra_body has a sampling_params_list field
-                combined_extra_body: dict[str, Any] = {"sampling_params_list": sampling_params}
+                # AR format: the payload has a sampling_params_list field, containing a list.
+                sampling_params_list = sampling_params if isinstance(sampling_params, list) else [sampling_params]
+                payload["sampling_params_list"] = sampling_params_list
 
         if negative_prompt:
             combined_extra_body["negative_prompt"] = negative_prompt
@@ -458,12 +445,11 @@ class VLLMOmniClient:
         if extra_body:
             combined_extra_body.update(extra_body)
 
-        payload: dict[str, Any] = {"messages": messages, "model": model}
+        # Add extra_body only if it has any content.
         if combined_extra_body:
             payload["extra_body"] = combined_extra_body
-        if modalities:
-            payload["modalities"] = modalities
 
+        # Place to inject any model-specific payload adjustment
         spec, _ = lookup_model_spec(model)
         if spec:
             preprocessor = spec.get("payload_preprocessor", None)

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/models.py

@@ -11,8 +11,7 @@ def _bagel_payload_preprocessor(payload: dict) -> dict:
                     content["text"] = "<|im_start|>" + content["text"] + "<|im_end|>"
     except (KeyError, TypeError):
         raise RuntimeError("Internal Error: malformatted BAGEL payload")
-    extra_body = payload.pop("extra_body", {})
-    return {**payload, **extra_body}
+    return payload
 
 
 def _qwen25_payload_preprocessor(payload: dict) -> dict:
@@ -37,7 +36,7 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
         ],
         "modes": [
             {
-                "mode": ModelMode.COMPREHENSION,
+                "mode": ModelMode.UNDERSTANDING,
                 "input_modalities": [Modality.TEXT, Modality.IMAGE],
             }
         ],
@@ -48,7 +47,7 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
         "payload_preprocessor": _qwen25_payload_preprocessor,
         "modes": [
             {
-                "mode": ModelMode.COMPREHENSION,
+                "mode": ModelMode.UNDERSTANDING,
                 "input_modalities": [
                     Modality.TEXT,
                     Modality.IMAGE,
@@ -62,7 +61,7 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
         "stages": ["autoregression", "autoregression", "autoregression"],
         "modes": [
             {
-                "mode": ModelMode.COMPREHENSION,
+                "mode": ModelMode.UNDERSTANDING,
                 "input_modalities": [
                     Modality.TEXT,
                     Modality.IMAGE,

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py

@@ -11,11 +11,23 @@ from typing import (
 AudioFormat: TypeAlias = Literal["mp3", "opus", "aac", "flac", "wav", "pcm"]
 
 
+class AutoregressionSamplingParams(dict):
+    pass
+
+
+class DiffusionSamplingParams(dict):
+    pass
+
+
+class QwenTTSModelSpecificParams(dict):
+    pass
+
+
 class ModelMode(Enum):
     IMAGE_GENERATION = auto()
     VIDEO_GENERATION = auto()
     AUDIO_GENERATION = auto()
-    COMPREHENSION = auto()
+    UNDERSTANDING = auto()
 
 
 class Modality(Enum):

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py

@@ -1,5 +1,6 @@
 from .logger import get_logger
 from .models import lookup_model_spec
+from .types import AutoregressionSamplingParams, DiffusionSamplingParams
 
 logger = get_logger(__name__)
 
@@ -31,24 +32,20 @@ def validate_model_and_sampling_params_types(
             )
         # Check that each stage's type match
         for i, sp in enumerate(sampling_param_list):
-            if "type" not in sp:
-                raise RuntimeError("Internal error: unknown sampling parameter type")
-            if sp["type"] != stages[i]:
+            if not _check_sampling_param_matches_stage(sp, stages[i]):
                 raise ValueError(
-                    f"Sampling parameter type ({sp['type']}) does not match "
+                    f"Sampling parameter type ({sp.__class__.__name__}) does not match "
                     f"stage type ({stages[i]}) at index {i} for model {model_name}."
                 )
     elif isinstance(sampling_param_list, dict):
-        if "type" not in sampling_param_list:
-            raise RuntimeError("Internal error: unknown sampling parameter type")
         # Check that the provided single sampling param matches all stages
-        elif any(stage != sampling_param_list["type"] for stage in stages):
-            raise ValueError(
-                f"When passing a single sampling parameter node, all stages of the model must match "
-                f"the provided sampling parameter's type. "
-                f"However, the stages of model {model_name} are: {stages}. "
-                f"The provided sampling parameter is {sampling_param_list['type']}"
-            )
+        for i, stage in enumerate(stages):
+            if not _check_sampling_param_matches_stage(sampling_param_list, stage):
+                raise ValueError(
+                    f"Provided single sampling parameter type ({sampling_param_list.__class__.__name__}) must match "
+                    f"the types of all stages of the model. "
+                    f"However, stage {i} of model {model_name} is of type {stage}."
+                )
 
 
 def add_sampling_parameters_to_stage(
@@ -82,3 +79,11 @@ def add_sampling_parameters_to_stage(
                 sampling_param_list[i].update(params_to_add)
 
     return sampling_param_list
+
+
+def _check_sampling_param_matches_stage(sampling_param: dict, stage_type: str) -> bool:
+    if stage_type == "autoregression":
+        return isinstance(sampling_param, AutoregressionSamplingParams)
+    if stage_type == "diffusion":
+        return isinstance(sampling_param, DiffusionSamplingParams)
+    raise RuntimeError(f"Internal error: unknown stage type {stage_type}.")

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Annotated Example.json

@@ -1 +0,0 @@
-{"id":"1c99f525-0a37-45ba-a28a-7df7c3af66b4","revision":0,"last_node_id":12,"last_link_id":14,"nodes":[{"id":1,"type":"VLLMOmniComprehension","pos":[1191.2177053682556,144.66829928181377],"size":[400,268],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":1},{"localized_name":"video","name":"video","shape":7,"type":"VIDEO","link":2},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":3},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":14},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"output_text","name":"output_text","type":"BOOLEAN","widget":{"name":"output_text"},"link":null},{"localized_name":"output_audio","name":"output_audio","type":"BOOLEAN","widget":{"name":"output_audio"},"link":null},{"localized_name":"use_audio_in_video","name":"use_audio_in_video","type":"BOOLEAN","widget":{"name":"use_audio_in_video"},"link":null}],"outputs":[{"localized_name":"text_response","name":"text_response","type":"STRING","links":[8]},{"localized_name":"audio_response","name":"audio_response","type":"AUDIO","links":[9]}],"properties":{"Node name for S&R":"VLLMOmniComprehension"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen2.5-Omni-7B","",true,true,true]},{"id":3,"type":"LoadVideo","pos":[729.5984141255855,-198.631920454299],"size":[282.798828125,233.0743408203125],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"file","name":"file","type":"COMBO","widget":{"name":"file"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[2]}],"properties":{"Node name for S&R":"LoadVideo"},"widgets_values":["draw.mp4","image"]},{"id":4,"type":"LoadAudio","pos":[729.8037086965753,99.86963519703949],"size":[282.798828125,136],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":5,"type":"VLLMOmniARSampling","pos":[510.3517536642828,658.073751009259],"size":[270,178],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[12]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":7,"type":"VLLMOmniARSampling","pos":[503.33235181647115,419.34158016181806],"size":[270,178],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[5,13]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":8,"type":"VLLMOmniSamplingParamsList","pos":[820.6056617389042,426.38372037182273],"size":[263.066015625,66],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"param1","name":"param1","type":"SAMPLING_PARAMS","link":5},{"localized_name":"param2","name":"param2","shape":7,"type":"SAMPLING_PARAMS","link":12},{"localized_name":"param3","name":"param3","shape":7,"type":"SAMPLING_PARAMS","link":13}],"outputs":[{"localized_name":"param list","name":"param list"Line truncated

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Chaining Services.json

@@ -0,0 +1,552 @@
+{
+    "id": "6643e5fd-fa2a-4f25-935a-483173a8097c",
+    "revision": 0,
+    "last_node_id": 8,
+    "last_link_id": 6,
+    "nodes": [
+        {
+            "id": 1,
+            "type": "PreviewImage",
+            "pos": [
+                1446.2005205859512,
+                -316.7359686049902
+            ],
+            "size": [
+                305.6628685610174,
+                307.0336884172169
+            ],
+            "flags": {},
+            "order": 5,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "images",
+                    "name": "images",
+                    "type": "IMAGE",
+                    "link": 1
+                }
+            ],
+            "outputs": [],
+            "properties": {
+                "Node name for S&R": "PreviewImage"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 8,
+            "type": "PreviewImage",
+            "pos": [
+                1010.1848219273943,
+                -311.54216706114363
+            ],
+            "size": [
+                305.6628685610174,
+                307.0336884172169
+            ],
+            "flags": {},
+            "order": 4,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "images",
+                    "name": "images",
+                    "type": "IMAGE",
+                    "link": 6
+                }
+            ],
+            "outputs": [],
+            "properties": {
+                "Node name for S&R": "PreviewImage"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 7,
+            "type": "VLLMOmniDiffusionSampling",
+            "pos": [
+                234.61958293819296,
+                9.405681270043619
+            ],
+            "size": [
+                284.205078125,
+                226
+            ],
+            "flags": {},
+            "order": 0,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "n",
+                    "name": "n",
+                    "type": "INT",
+                    "widget": {
+                        "name": "n"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "num_inference_steps",
+                    "name": "num_inference_steps",
+                    "type": "INT",
+                    "widget": {
+                        "name": "num_inference_steps"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "guidance_scale",
+                    "name": "guidance_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "guidance_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "true_cfg_scale",
+                    "name": "true_cfg_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "true_cfg_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_slicing",
+                    "name": "vae_use_slicing",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_slicing"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_tiling",
+                    "name": "vae_use_tiling",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_tiling"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "diffusion sampling params",
+                    "name": "diffusion sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        5
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniDiffusionSampling"
+            },
+            "widgets_values": [
+                1,
+                50,
+                1,
+                1,
+                false,
+                false,
+                1525,
+                "randomize"
+            ]
+        },
+        {
+            "id": 4,
+            "type": "VLLMOmniDiffusionSampling",
+            "pos": [
+                666.8380026154548,
+                268.86271068330126
+            ],
+            "size": [
+                284.205078125,
+                226
+            ],
+            "flags": {},
+            "order": 1,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "n",
+                    "name": "n",
+                    "type": "INT",
+                    "widget": {
+                        "name": "n"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "num_inference_steps",
+                    "name": "num_inference_steps",
+                    "type": "INT",
+                    "widget": {
+                        "name": "num_inference_steps"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "guidance_scale",
+                    "name": "guidance_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "guidance_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "true_cfg_scale",
+                    "name": "true_cfg_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "true_cfg_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_slicing",
+                    "name": "vae_use_slicing",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_slicing"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_tiling",
+                    "name": "vae_use_tiling",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_tiling"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "diffusion sampling params",
+                    "name": "diffusion sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        2
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniDiffusionSampling"
+            },
+            "widgets_values": [
+                4,
+                50,
+                7,
+                1,
+                false,
+                false,
+                42,
+                "fixed"
+            ]
+        },
+        {
+            "id": 5,
+            "type": "VLLMOmniGenerateImage",
+            "pos": [
+                984.723613585788,
+                63.376900027553276
+            ],
+            "size": [
+                416.56628685610167,
+                372.1662621294205
+            ],
+            "flags": {},
+            "order": 3,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": 4
+                },
+                {
+                    "localized_name": "mask",
+                    "name": "mask",
+                    "shape": 7,
+                    "type": "MASK",
+                    "link": null
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 2
+                },
+                {
+                    "localized_name": "url",
+                    "name": "url",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "url"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "model",
+                    "name": "model",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "model"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "prompt",
+                    "name": "prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "negative_prompt",
+                    "name": "negative_prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "negative_prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "width",
+                    "name": "width",
+                    "type": "INT",
+                    "widget": {
+                        "name": "width"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "height",
+                    "name": "height",
+                    "type": "INT",
+                    "widget": {
+                        "name": "height"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "type": "IMAGE",
+                    "links": [
+                        1
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniGenerateImage"
+            },
+            "widgets_values": [
+                "http://localhost:8001/v1",
+                "/home/models/Qwen/Qwen-Image-Edit",
+                "A high-quality, high contrast, stylized portrait of the object in the uploaded reference image. Pop art and doodle style, with abundant scribbling patterns, such as a teal crown, orange lightning bolts, colorful handwritten scripts.",
+                "Realistic",
+                800,
+                800
+            ]
+        },
+        {
+            "id": 6,
+            "type": "VLLMOmniGenerateImage",
+            "pos": [
+                541.2702717429818,
+                -160.70392744708548
+            ],
+            "size": [
+                416.56628685610167,
+                372.1662621294205
+            ],
+            "flags": {},
+            "order": 2,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": null
+                },
+                {
+                    "localized_name": "mask",
+                    "name": "mask",
+                    "shape": 7,
+                    "type": "MASK",
+                    "link": null
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 5
+                },
+                {
+                    "localized_name": "url",
+                    "name": "url",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "url"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "model",
+                    "name": "model",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "model"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "prompt",
+                    "name": "prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "negative_prompt",
+                    "name": "negative_prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "negative_prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "width",
+                    "name": "width",
+                    "type": "INT",
+                    "widget": {
+                        "name": "width"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "height",
+                    "name": "height",
+                    "type": "INT",
+                    "widget": {
+                        "name": "height"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "type": "IMAGE",
+                    "links": [
+                        4,
+                        6
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniGenerateImage"
+            },
+            "widgets_values": [
+                "http://localhost:8000/v1",
+                "/home/models/Tongyi-MAI/Z-Image-Turbo",
+                "A headshot of a cute Siamese kitty. Blurred background due to wide aperture. Close-up look. Realistic.",
+                "Cartoonish.",
+                800,
+                800
+            ]
+        }
+    ],
+    "links": [
+        [
+            1,
+            5,
+            0,
+            1,
+            0,
+            "IMAGE"
+        ],
+        [
+            2,
+            4,
+            0,
+            5,
+            2,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            4,
+            6,
+            0,
+            5,
+            0,
+            "IMAGE"
+        ],
+        [
+            5,
+            7,
+            0,
+            6,
+            2,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            6,
+            6,
+            0,
+            8,
+            0,
+            "IMAGE"
+        ]
+    ],
+    "groups": [],
+    "config": {},
+    "extra": {
+        "workflowRendererVersion": "LG",
+        "ds": {
+            "scale": 1.0426432563169903,
+            "offset": [
+                -90.4071374799346,
+                575.2948920069742
+            ]
+        }
+    },
+    "version": 0.4
+}

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Image Generation.json

@@ -1 +1 @@
-{"id":"91f75acc-8040-40f6-865a-2e8a7cfd6672","revision":0,"last_node_id":11,"last_link_id":23,"nodes":[{"id":3,"type":"PreviewImage","pos":[1281.8455167767304,-69.02638461454333],"size":[305.6628685610174,307.0336884172169],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":22}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"VLLMOmniGenerateImage","pos":[816.0962515873642,-137.51112584854445],"size":[400,278],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":21},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":23},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[22]}],"properties":{"Node name for S&R":"VLLMOmniGenerateImage"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen-Image-Edit","Put this figure in a realistic mountain view","",512,512]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,202],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scale"},"link":null},{"localized_name":"vae_use_slicing","name":"vae_use_slicing","type":"BOOLEAN","widget":{"name":"vae_use_slicing"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[23]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[4,50,1,1,false,42,"fixed"]},{"id":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Image\" node supports both text-to-image generation or image-to-image generation (a.k.a. image editing). The node will route to the correct endpoint depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[496.31859627609606,-229.71277089860084],"size":[270,314],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[21]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]}],"links":[[21,4,0,11,0,"IMAGE"],[22,11,0,3,0,"IMAGE"],[23,8,0,11,2,"SAMPLING_PARAMS"]],"groups":[{"id":1,"title":"Input","bounding":[213.1706010087147,-313.20095750667554,560.8246471437027,407.4246532472182],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486894,"offset":[276.8427259980847,549.4885226647582]}},"version":0.4}
+{"id":"91f75acc-8040-40f6-865a-2e8a7cfd6672","revision":0,"last_node_id":11,"last_link_id":23,"nodes":[{"id":3,"type":"PreviewImage","pos":[1281.8455167767304,-69.02638461454333],"size":[305.6628685610174,307.0336884172169],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"图像","name":"images","type":"IMAGE","link":22}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"VLLMOmniGenerateImage","pos":[816.0962515873642,-137.51112584854445],"size":[400,278],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":21},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":23},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[22]}],"properties":{"Node name for S&R":"VLLMOmniGenerateImage"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen-Image-Edit","Put this figure in a realistic mountain view","",512,512]},{"id":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Image\" node supports both text-to-image generation or image-to-image generation (a.k.a. image editing). The node will route to the correct endpoint depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[496.31859627609606,-229.71277089860084],"size":[270,314],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"图像","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"选择文件上传","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"图像","name":"IMAGE","type":"IMAGE","links":[21]},{"localized_name":"遮罩","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,226],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scale"},"link":null},{"localized_name":"vae_use_slicing","name":"vae_use_slicing","type":"BOOLEAN","widget":{"name":"vae_use_slicing"},"link":null},{"localized_name":"vae_use_tiling","name":"vae_use_tiling","type":"BOOLEAN","widget":{"name":"vae_use_tiling"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[23]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[4,50,1,1,false,false,42,"randomize"]}],"links":[[21,4,0,11,0,"IMAGE"],[22,11,0,3,0,"IMAGE"],[23,8,0,11,2,"SAMPLING_PARAMS"]],"groups":[{"id":1,"title":"Input","bounding":[213.1706010087147,-313.20095750667554,560.8246471437027,407.4246532472182],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486894,"offset":[-55.18047513099182,220.2553505195962]}},"version":0.4}

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Multimodal Understanding.json

@@ -0,0 +1,761 @@
+{
+    "id": "1c99f525-0a37-45ba-a28a-7df7c3af66b4",
+    "revision": 0,
+    "last_node_id": 12,
+    "last_link_id": 14,
+    "nodes": [
+        {
+            "id": 1,
+            "type": "VLLMOmniUnderstanding",
+            "pos": [
+                1191.2177053682556,
+                144.66829928181377
+            ],
+            "size": [
+                400,
+                268
+            ],
+            "flags": {},
+            "order": 8,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": 1
+                },
+                {
+                    "localized_name": "video",
+                    "name": "video",
+                    "shape": 7,
+                    "type": "VIDEO",
+                    "link": 2
+                },
+                {
+                    "localized_name": "audio",
+                    "name": "audio",
+                    "shape": 7,
+                    "type": "AUDIO",
+                    "link": 3
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 14
+                },
+                {
+                    "localized_name": "url",
+                    "name": "url",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "url"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "model",
+                    "name": "model",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "model"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "prompt",
+                    "name": "prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "output_text",
+                    "name": "output_text",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "output_text"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "output_audio",
+                    "name": "output_audio",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "output_audio"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "use_audio_in_video",
+                    "name": "use_audio_in_video",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "use_audio_in_video"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "text_response",
+                    "name": "text_response",
+                    "type": "STRING",
+                    "links": [
+                        8
+                    ]
+                },
+                {
+                    "localized_name": "audio_response",
+                    "name": "audio_response",
+                    "type": "AUDIO",
+                    "links": [
+                        9
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniUnderstanding"
+            },
+            "widgets_values": [
+                "http://localhost:8000/v1",
+                "Qwen/Qwen2.5-Omni-7B",
+                "",
+                true,
+                true,
+                true
+            ]
+        },
+        {
+            "id": 3,
+            "type": "LoadVideo",
+            "pos": [
+                729.5984141255855,
+                -198.631920454299
+            ],
+            "size": [
+                282.798828125,
+                233.0743408203125
+            ],
+            "flags": {},
+            "order": 0,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "file",
+                    "name": "file",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "file"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "choose file to upload",
+                    "name": "upload",
+                    "type": "IMAGEUPLOAD",
+                    "widget": {
+                        "name": "upload"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "VIDEO",
+                    "name": "VIDEO",
+                    "type": "VIDEO",
+                    "links": [
+                        2
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadVideo"
+            },
+            "widgets_values": [
+                "draw.mp4",
+                "image"
+            ]
+        },
+        {
+            "id": 4,
+            "type": "LoadAudio",
+            "pos": [
+                729.8037086965753,
+                99.86963519703949
+            ],
+            "size": [
+                282.798828125,
+                136
+            ],
+            "flags": {},
+            "order": 1,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "audio",
+                    "name": "audio",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "audio"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "audioUI",
+                    "name": "audioUI",
+                    "type": "AUDIO_UI",
+                    "widget": {
+                        "name": "audioUI"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "choose file to upload",
+                    "name": "upload",
+                    "type": "AUDIOUPLOAD",
+                    "widget": {
+                        "name": "upload"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "AUDIO",
+                    "name": "AUDIO",
+                    "type": "AUDIO",
+                    "links": [
+                        3
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadAudio"
+            },
+            "widgets_values": [
+                "Megan-Fox.mp3",
+                null,
+                null
+            ]
+        },
+        {
+            "id": 5,
+            "type": "VLLMOmniARSampling",
+            "pos": [
+                510.3517536642828,
+                658.073751009259
+            ],
+            "size": [
+                270,
+                178
+            ],
+            "flags": {},
+            "order": 2,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "max_tokens",
+                    "name": "max_tokens",
+                    "type": "INT",
+                    "widget": {
+                        "name": "max_tokens"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "temperature",
+                    "name": "temperature",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "temperature"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "top_p",
+                    "name": "top_p",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "top_p"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "repetition_penalty",
+                    "name": "repetition_penalty",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "repetition_penalty"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "AR sampling params",
+                    "name": "AR sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        12
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniARSampling"
+            },
+            "widgets_values": [
+                100,
+                1,
+                1,
+                1,
+                -1,
+                "randomize"
+            ]
+        },
+        {
+            "id": 7,
+            "type": "VLLMOmniARSampling",
+            "pos": [
+                503.33235181647115,
+                419.34158016181806
+            ],
+            "size": [
+                270,
+                178
+            ],
+            "flags": {},
+            "order": 3,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "max_tokens",
+                    "name": "max_tokens",
+                    "type": "INT",
+                    "widget": {
+                        "name": "max_tokens"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "temperature",
+                    "name": "temperature",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "temperature"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "top_p",
+                    "name": "top_p",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "top_p"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "repetition_penalty",
+                    "name": "repetition_penalty",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "repetition_penalty"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "AR sampling params",
+                    "name": "AR sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        5,
+                        13
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniARSampling"
+            },
+            "widgets_values": [
+                100,
+                1,
+                1,
+                1,
+                -1,
+                "randomize"
+            ]
+        },
+        {
+            "id": 8,
+            "type": "VLLMOmniSamplingParamsList",
+            "pos": [
+                820.6056617389042,
+                426.38372037182273
+            ],
+            "size": [
+                263.066015625,
+                66
+            ],
+            "flags": {},
+            "order": 7,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "param1",
+                    "name": "param1",
+                    "type": "SAMPLING_PARAMS",
+                    "link": 5
+                },
+                {
+                    "localized_name": "param2",
+                    "name": "param2",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 12
+                },
+                {
+                    "localized_name": "param3",
+                    "name": "param3",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 13
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "param list",
+                    "name": "param list",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        14
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniSamplingParamsList"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 11,
+            "type": "MarkdownNote",
+            "pos": [
+                826.2328280438272,
+                569.1890318701705
+            ],
+            "size": [
+                333.8220435590464,
+                261.63596728060656
+            ],
+            "flags": {},
+            "order": 4,
+            "mode": 0,
+            "inputs": [],
+            "outputs": [],
+            "title": "Note: Sampling Parameters",
+            "properties": {},
+            "widgets_values": [
+                "## Sampling Parameter Types\n\nThere are two types of sampling parameters: one for autoregression and one for diffusion.\nYou should ensure that you have chosen the correct type of sampling parameters for the model you request.\n\n## Stages & Shorthand\n\nFor multi-stage models such as Qwen Omni, you can either\n- connect one sampling parameter node, which is applied to all stages.\n- connect exactly the same number of sampling parameter nodes to a \"Multi-Stage Sampling Parameter List\" node, then connect this node to the primary request node.\n\nNote that this shorthand is intended to stay consistent with the [online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/qwen2_5_omni/)"
+            ],
+            "color": "#432",
+            "bgcolor": "#000"
+        },
+        {
+            "id": 12,
+            "type": "MarkdownNote",
+            "pos": [
+                378.9866207003777,
+                152.59550252215752
+            ],
+            "size": [
+                319.7287574247016,
+                107.15904081906785
+            ],
+            "flags": {},
+            "order": 6,
+            "mode": 0,
+            "inputs": [],
+            "outputs": [],
+            "title": "Note: Input",
+            "properties": {},
+            "widgets_values": [
+                "Note that not all models support every modality as input. For example, `ByteDance-Seed/BAGEL-7B-MoT` in Multimodality Understanding mode only support text and image input.\n\nYou should ensure that the input are supported by the model. You can check the corresponding [online serving documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/) for confirmation."
+            ],
+            "color": "#432",
+            "bgcolor": "#000"
+        },
+        {
+            "id": 2,
+            "type": "LoadImage",
+            "pos": [
+                394.4674804308822,
+                -207.6987397548834
+            ],
+            "size": [
+                282.798828125,
+                314
+            ],
+            "flags": {},
+            "order": 5,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "image"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "choose file to upload",
+                    "name": "upload",
+                    "type": "IMAGEUPLOAD",
+                    "widget": {
+                        "name": "upload"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "IMAGE",
+                    "name": "IMAGE",
+                    "type": "IMAGE",
+                    "links": [
+                        1
+                    ]
+                },
+                {
+                    "localized_name": "MASK",
+                    "name": "MASK",
+                    "type": "MASK",
+                    "links": null
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadImage"
+            },
+            "widgets_values": [
+                "example.png",
+                "image"
+            ]
+        },
+        {
+            "id": 10,
+            "type": "PreviewAudio",
+            "pos": [
+                1664.548345556043,
+                297.5921292054054
+            ],
+            "size": [
+                270,
+                88
+            ],
+            "flags": {},
+            "order": 10,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "audio",
+                    "name": "audio",
+                    "type": "AUDIO",
+                    "link": 9
+                },
+                {
+                    "localized_name": "audioUI",
+                    "name": "audioUI",
+                    "type": "AUDIO_UI",
+                    "widget": {
+                        "name": "audioUI"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [],
+            "properties": {
+                "Node name for S&R": "PreviewAudio"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 9,
+            "type": "ShowText|pysssss",
+            "pos": [
+                1649.2506875091847,
+                66.22823888292349
+            ],
+            "size": [
+                318.7188464232943,
+                173.38502269972975
+            ],
+            "flags": {},
+            "order": 9,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "text",
+                    "name": "text",
+                    "type": "STRING",
+                    "link": 8
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "STRING",
+                    "name": "STRING",
+                    "shape": 6,
+                    "type": "STRING",
+                    "links": null
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "ShowText|pysssss"
+            },
+            "widgets_values": []
+        }
+    ],
+    "links": [
+        [
+            1,
+            2,
+            0,
+            1,
+            0,
+            "IMAGE"
+        ],
+        [
+            2,
+            3,
+            0,
+            1,
+            1,
+            "VIDEO"
+        ],
+        [
+            3,
+            4,
+            0,
+            1,
+            2,
+            "AUDIO"
+        ],
+        [
+            5,
+            7,
+            0,
+            8,
+            0,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            8,
+            1,
+            0,
+            9,
+            0,
+            "STRING"
+        ],
+        [
+            9,
+            1,
+            1,
+            10,
+            0,
+            "AUDIO"
+        ],
+        [
+            12,
+            5,
+            0,
+            8,
+            1,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            13,
+            7,
+            0,
+            8,
+            2,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            14,
+            8,
+            0,
+            1,
+            3,
+            "SAMPLING_PARAMS"
+        ]
+    ],
+    "groups": [
+        {
+            "id": 1,
+            "title": "Sampling Parameters",
+            "bounding": [
+                480.1649301181556,
+                341.08402513937995,
+                692.4113568972277,
+                510.48648853403665
+            ],
+            "color": "#3f789e",
+            "font_size": 24,
+            "flags": {}
+        },
+        {
+            "id": 2,
+            "title": "Input",
+            "bounding": [
+                344.2364817813528,
+                -287.5850484183313,
+                704.8832238768634,
+                559.5124009832894
+            ],
+            "color": "#3f789e",
+            "font_size": 24,
+            "flags": {}
+        }
+    ],
+    "config": {},
+    "extra": {
+        "workflowRendererVersion": "LG",
+        "ds": {
+            "scale": 0.9478575057427204,
+            "offset": [
+                33.81037136029557,
+                307.1974296197726
+            ]
+        }
+    },
+    "version": 0.4
+}

docker/Dockerfile.ci

@@ -8,7 +8,7 @@ COPY . .
 
 # Install system dependencies
 RUN apt-get update && \
-    apt-get install -y ffmpeg sox libsox-fmt-all jq && \
+    apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 

docker/Dockerfile.npu

@@ -2,13 +2,15 @@ ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
 ARG VLLM_ASCEND_TAG=v0.14.0rc1
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
+WORKDIR /vllm-workspace/vllm-ascend
+RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+RUN pip install -v -e .
+
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 
 COPY . .
 
-RUN sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
- && sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
 RUN pip install -v -e .
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn

docker/Dockerfile.npu.a3

@@ -2,13 +2,15 @@ ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
 ARG VLLM_ASCEND_TAG=v0.14.0rc1-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
+WORKDIR /vllm-workspace/vllm-ascend
+RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+RUN pip install -v -e .
+
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 
 COPY . .
 
-RUN sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
- && sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
 RUN pip install -v -e .
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn

docker/Dockerfile.rocm

@@ -19,7 +19,7 @@ WORKDIR ${COMMON_WORKDIR}
 
 # Step 1: Setup - Install system dependencies
 RUN apt-get update && \
-    apt-get install -y ffmpeg && \
+    apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 

docs/.nav.yml

@@ -13,9 +13,10 @@ nav:
     - examples/README.md
     - Offline Inference:
       - BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
+      - GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
       - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
-      - LoRA Inference(Diffusion): user_guide/examples/offline_inference/lora_inference.md
+      - MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
       - Qwen3-TTS: user_guide/examples/offline_inference/qwen3_tts.md
@@ -24,12 +25,15 @@ nav:
       - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
     - Online Serving:
       - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
+      - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
       - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
-      - LoRA Inference(Diffusion): user_guide/examples/online_serving/lora_inference.md
+      - Image-To-Video: user_guide/examples/online_serving/image_to_video.md
+      - Online serving Example of vLLM-Omni for MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
       - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
       - Qwen3-TTS: user_guide/examples/online_serving/qwen3_tts.md
       - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
+      - Text-To-Video: user_guide/examples/online_serving/text_to_video.md
   - General:
     - usage/*
   - Configuration:
@@ -46,13 +50,17 @@ nav:
       - Quantization:
         - Overview: user_guide/diffusion/quantization/overview.md
         - FP8: user_guide/diffusion/quantization/fp8.md
+        - GGUF: user_guide/diffusion/quantization/gguf.md
       - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
       - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
       - LoRA: user_guide/diffusion/lora.md
+      - Hybrid Sharded Data Parallel: design/feature/hsdp.md
+      - Custom Pipeline: features/custom_pipeline.md
     - ComfyUI: features/comfyui.md
 - Developer Guide:
   - General:
     - contributing/README.md
+    - pr_reviewer.md
     - glob: contributing/*
       flatten_single_child_sections: true
   - Model Implementation:
@@ -72,6 +80,7 @@ nav:
       - design/feature/tensor_parallel.md
       - design/feature/cache_dit.md
       - design/feature/teacache.md
+      - design/feature/async_chunk_design.md
     - Module Design:
       - design/module/ar_module.md
       - design/module/dit_module.md

docs/api/README.md

@@ -7,7 +7,6 @@ Main entry points for vLLM-Omni inference and serving.
 - [vllm_omni.entrypoints.async_omni.AsyncOmni][]
 - [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
 - [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
-- [vllm_omni.entrypoints.chat_utils.extract_audio_from_video_async][]
 - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
 - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
 - [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
@@ -19,6 +18,7 @@ Main entry points for vLLM-Omni inference and serving.
 - [vllm_omni.entrypoints.omni_llm.OmniLLM][]
 - [vllm_omni.entrypoints.omni_stage.OmniStage][]
 - [vllm_omni.entrypoints.stage_utils.OmniStageTaskType][]
+- [vllm_omni.entrypoints.zmq_utils.ZmqQueue][]
 
 ## Inputs
 
@@ -36,6 +36,12 @@ Input data structures for multi-modal inputs.
 Engine classes for offline and online inference.
 
 - [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.BufferAllocator][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.ManagedBuffer][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.MooncakeAgentMetadata][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.MooncakeTransferEngineConnector][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.QueryRequest][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.QueryResponse][]
 - [vllm_omni.engine.AdditionalInformationEntry][]
 - [vllm_omni.engine.AdditionalInformationPayload][]
 - [vllm_omni.engine.OmniEngineCoreOutput][]
@@ -72,6 +78,8 @@ Configuration classes.
 - [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
 - [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][]
 - [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][]
+- [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioConfig][]
+- [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioTokenizerConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSSpeakerEncoderConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSTalkerCodePredictorConfig][]
@@ -89,8 +97,10 @@ Configuration classes.
 Worker classes and model runners for distributed inference.
 
 - [vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner][]
+- [vllm_omni.diffusion.worker.diffusion_worker.CustomPipelineWorkerExtension][]
 - [vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker][]
 - [vllm_omni.diffusion.worker.diffusion_worker.WorkerProc][]
+- [vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase][]
 - [vllm_omni.platforms.npu.worker.npu_ar_model_runner.ExecuteModelState][]
 - [vllm_omni.platforms.npu.worker.npu_ar_model_runner.NPUARModelRunner][]
 - [vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker][]

docs/contributing/profiling.md

@@ -131,9 +131,6 @@ python image_to_video.py \
 
 2. **Wan-AI/Wan2.2-I2V-A14B-Diffusers**:   [https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
 
-> **Note:**
-As of now, asynchronous (online) profiling is not fully supported in vLLM-Omni. While start_profile() and stop_profile() methods exist, they are only reliable in offline inference scripts (e.g., the provided end2end.py examples). Do not use them in server-mode or streaming scenarios—traces may be incomplete or fail to flush.
-
 ### 4. Analyzing Omni Traces
 
 Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.

docs/design/feature/async_chunk_design.md

@@ -0,0 +1,209 @@
+# Async Chunk Design
+
+## Table of Contents
+
+1. [Overview](#overview)
+2. [Performance](#performance)
+3. [Architecture](#architecture)
+4. [Configuration](#configuration)
+5. [Related Files](#related-files)
+
+## Overview
+
+The `async_chunk` feature enables asynchronous, chunked processing of data across multiple stages in a multi-stage pipeline (e.g., Qwen3-Omni with Thinker → Talker → Code2Wav stages). Instead of waiting for a complete stage output before forwarding to the next stage, this feature allows stages to process and forward data in chunks as it becomes available, significantly reducing latency and improving throughput.
+
+**Chunk Size Definition**
+
+- **Prefill Phase**: `chunk_size = num_scheduled_tokens` for chunked prefill processing
+- **Decode Phase**: `chunk_size = num_scheduled_tokens = 1 ` for per-token streaming
+
+For qwen3-omni:
+- **Thinker → Talker**: Per decode step (typically chunk_size=1)
+- **Talker → Code2Wav**: Accumulated to code2wav chunk_size(default=25, currently only support default, will support chunk_size soon) before sending
+- **Code2Wav**: Streaming decode with code2wav chunk_size
+
+With `async_chunk`:
+- Stages can start processing as soon as chunks are available
+- Overlapping execution across stages
+- Reduced latency and improved throughput
+- Better resource utilization
+- Async scheduling: Chunk IO (get/put) overlaps with compute via background threads so the scheduler is not blocked waiting for chunks
+
+## Performance
+1. **Reduced Latency**: Next stage can start processing immediately
+2. **Streaming Support**: Enables streaming for audio generation
+3. **IO-Compute Overlap**: Chunk retrieval happens asynchronously while other requests compute
+4. **Non-blocking Scheduler**: Requests waiting for chunks don't block the entire scheduler
+5. **Code2Wav Batch Inference**: Supports batched processing in code2wav stage
+
+| Input     | Output           | Async_chunk enabled | Code2Wav batch size | Max_Concurrency | Prompts | Mean E2E   | Mean TTFT  | Mean TPOT | Mean TTFP   | Mean RTF | Mean ITL |
+|-----------|------------------|---------------------|---------------------|----------------|---------|------------|------------|-----------|-------------|----------|----------|
+| text 100  | text 100+audio   | False               | 1                   | 1              | 50      | 6581.80    | 43.22      | 8.31      | 6459.34     | 0.24     | 8.22     |
+| text 100  | text 100+audio   | False               | 1                   | 4              | 50      | 7398.63    | 67.57      | 9.14      | 7285.35     | 0.27     | 9.05     |
+| text 100  | text 100+audio   | False               | 1                   | 10             | 50      | 13522.99   | 131.82     | 12.72     | 13410.44    | 0.49     | 12.60    |
+| text 100  | text 100+audio   | False               | 64                  | 1              | 50      | 6505.13    | 43.14      | 8.52      | 6395.40     | 0.24     | 8.44     |
+| text 100  | text 100+audio   | False               | 64                  | 4              | 50      | 7668.15    | 51.15      | 9.36      | 7562.37     | 0.28     | 9.27     |
+| text 100  | text 100+audio   | False               | 64                  | 10             | 50      | 9516.18    | 138.06     | 14.75     | 9409.26     | 0.34     | 14.60    |
+| text 100  | text 100+audio   | True                | 1                   | 1              | 50      | 6179.79    | 44.58      | 8.69      | 522.99      | 0.22     | 8.60     |
+| text 100  | text 100+audio   | True                | 1                   | 4              | 50      | 7692.69    | 103.96     | 10.22     | 785.85      | 0.29     | 10.12    |
+| text 100  | text 100+audio   | True                | 1                   | 10             | 50      | 11152.71   | 685.60     | 17.64     | 1628.88     | 0.41     | 17.62    |
+
+
+Performance data collected on H800 GPUs through comprehensive benchmarking with cudagraph enabled. text input uses random dataset.
+
+Enabling **async_chunk** (False→True) sharply reduces time-to-first-audio (TTFP)—e.g. ~92% at concurrency 1 (6.5s→0.52s)—and improves E2E latency (e.g. ~6% at conc 1, ~17% at conc 10). RTF (Real Time Factor) also improves with async_chunk on (e.g. ~8% at conc 1: 0.24→0.22, ~16% at conc 10: 0.49→0.41). Enabling **Code2Wav batch size 64** (vs 1) improves E2E and TTFP at higher concurrency when async_chunk is off (e.g. ~30% at conc 10: 13.5s→9.5s E2E, 13.4s→9.4s TTFP).
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_ttfp_performance.png">
+    <img alt="TTFP Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_ttfp_performance.png" width=100%>
+  </picture>
+</p>
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_rtf_performance.png">
+    <img alt="RTF Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_rtf_performance.png" width=100%>
+  </picture>
+</p>
+
+## Architecture
+### Data Flow
+
+#### Sequential Flow
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-non-async-chunk.png">
+    <img alt="Data Flow between stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-non-async-chunk.png" width=100%>
+  </picture>
+</p>
+
+#### Async Chunk Flow
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png">
+    <img alt="Data Flow between stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png" width=100%>
+  </picture>
+</p>
+
+### Async Chunk architecture
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/async-chunk-architecture.png">
+    <img alt="Async Chunk Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/async-chunk-architecture.png" width=100%>
+  </picture>
+</p>
+
+
+### Key Components
+
+1. **OmniConnector**: Inter-stage data transport only
+   - Shared memory or other IPC mechanisms
+   - **Transport-only API**: `put(from_stage, to_stage, put_key, data)` and `get(from_stage, to_stage, get_key)` (optionally with timeout)
+   - **No request-specific state**: Connector does not track put_requests, get_requests, request_payload, finished_requests, or other request-bound metadata; it only performs put/get operations
+   - Chunk keys and request/chunk lifecycle are managed by **OmniChunkTransferAdapter**
+
+2. **Transfer Adapter Layer**: Extensible abstraction for managing data transfer via connectors
+   - **OmniTransferAdapterBase**: Base class with background **recv_loop** and **save_loop** threads;
+   - **OmniChunkTransferAdapter**: Chunk-specific implementation that owns the full chunk lifecycle when async_chunk is enabled
+     - **Chunk ID and key construction**: Builds keys like `{req_id}_{stage_id}_{chunk_id}` for put/get
+     - **Async get**: `load_async(request)` enqueues the request; background **recv_loop** polls the connector (non-blocking); when data is available, updates the request and marks it in `_finished_load_reqs`; scheduler calls `get_finished_requests()` to learn which requests have chunks ready
+     - **Async put**: `save_async(pooling_output, request)` invokes `custom_process_next_stage_input_func` in the main thread to build the payload, then enqueues a save task; background **save_loop** performs `connector.put()`; payload processing and chunk accumulation (e.g. code2wav chunk_size) remain in the main thread
+
+3. **Stage Input Processors**: Custom functions that process stage outputs into chunks for different models
+   - Receive **transfer_manager** (OmniChunkTransferAdapter)
+   - Qwen3-omni reference: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
+
+4. **Schedulers**: Modified to handle chunk-based scheduling with async IO-compute overlap
+   - `OmniARScheduler`: For autoregressive stages
+   - `OmniGenerationScheduler`: For generation stages
+   - Both schedulers use **OmniChunkTransferAdapter** and **before/after** hooks around `super().schedule()`:
+     - **Before** `super().schedule()`: `process_pending_chunks(waiting, running)` moves requests waiting for chunks to `WAITING_FOR_CHUNK`, enqueues load tasks for background polling
+     - **After** `super().schedule()`: `restore_queues(waiting, running)` restores requests with ready chunks back to waiting/running, `postprocess_scheduler_output(scheduler_output)` attaches cached additional_information, clears chunk-ready flags
+   - **put_chunk** `save_async(pooler_output, request)`; **get_chunk** / **get_chunk_for_generation** `load_async(request)`
+
+5. **Model Runners**: Handle chunk processing
+   - `OmniGPUModelRunner`: Processes chunks in AR stages
+   - `GPUGenerationModelRunner`: Processes chunks in generation stages
+     - Uses `ubatch_slices` from `get_forward_context()` to track per-request sequence lengths in batched inference
+     - Reuses `ubatch_slices_padded` for code2wav batching to properly split batch outputs
+     - Handles list-type multimodal outputs: iterates through requests and assigns corresponding tensor to each
+     - Improved request state management: removes unscheduled and finished requests from input batch
+
+6. **Model Implementation**: Model-specific chunk handling
+   - `Qwen3OmniMoeForConditionalGeneration`: Main model with async_chunk support
+     - **Code2Wav stage batching**: Uses `ubatch_slices` to construct batched codec codes tensor `[batch_size, 16, max_seq_len]`
+     - **Batch output handling**: `generate_audio()` returns `list[torch.Tensor]`, one audio tensor per request
+     - **Multimodal outputs**: Returns list of audio tensors for batch processing instead of single concatenated tensor
+   - `Qwen3OmniCode2WavDecoder`: Audio generation model
+     - `chunked_decode()` and `chunked_decode_streaming()`: Return `list[torch.Tensor]` (one per request)
+     - Uses `ubatch_slices` to split batched waveform output into per-request audio chunks
+     - Each request gets correctly sized audio based on its code sequence length: `waveform[:, :, :code_seq_len * total_upsample]`
+
+7. **Request status**: `RequestStatus.WAITING_FOR_CHUNK` is added via patch (e.g. in `vllm_omni/patch.py`) so requests waiting for a chunk are not scheduled by the base vLLM scheduler until the chunk is ready.
+
+## Configuration
+
+Enable async_chunk in stage configuration YAML:
+
+```yaml
+async_chunk: true
+stage_args:
+  - stage_id: 0
+    engine_args:
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk
+  - stage_id: 1
+    engine_args:
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk
+```
+
+### Stage Configuration
+
+- `async_chunk: bool`: Enable/disable async chunk mode
+- `custom_process_next_stage_input_func: str`: Path to custom chunk processing function; receives `(transfer_manager, pooling_output, request)`. For qwen3-omni: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
+- `stage_connector_config: dict`: Connector configuration
+- `worker_type: str`: Model type, e.g. `"ar"` or `"generation"` (used by OmniChunkTransferAdapter for mode-specific payload handling)
+- `max_batch_size: int`: Maximum batch size for the stage
+
+
+### Connector Configuration
+
+```yaml
+connectors:
+  - from_stage: 0
+    to_stage: 1
+    spec:
+      name: SharedMemoryConnector
+      extra:
+        stage_id: 0
+```
+
+### Code2Wav Batch Configuration
+
+For optimal performance with async_chunk, the code2wav stage should be configured with batching:
+
+```yaml
+stage_args:
+  - stage_id: 2  # code2wav stage
+    runtime:
+      devices: "1"
+      max_batch_size: 64  # Enables batched audio generation
+    engine_args:
+      model_stage: code2wav
+```
+
+## Related Files
+
+- `vllm_omni/model_executor/stage_input_processors/qwen3_omni.py`: Chunk processing functions (receive `transfer_manager` as first param)
+- `vllm_omni/distributed/omni_connectors/transfer_adapter/base.py`: OmniTransferAdapterBase (recv_loop, save_loop, load_async, save_async)
+- `vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py`: OmniChunkTransferAdapter (process_pending_chunks, restore_queues, postprocess_scheduler_output)
+- `vllm_omni/distributed/omni_connectors/connectors/shm_connector.py`: SharedMemoryConnector (transport-only put/get)
+- `vllm_omni/core/sched/omni_ar_scheduler.py`: AR scheduler with chunk_transfer_adapter
+- `vllm_omni/core/sched/omni_generation_scheduler.py`: Generation scheduler with same async chunk pattern
+- `vllm_omni/worker/gpu_model_runner.py`: Model runner with chunk handling
+- `vllm_omni/worker/gpu_generation_model_runner.py`: Generation model runner with batch output handling and ubatch_slices support
+- `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py`: Model implementation with code2wav batching
+- `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py`: Code2wav decoder with batch support
+- `vllm_omni/engine/arg_utils.py`: Configuration definitions (async_chunk, worker_type)
+- `vllm_omni/config/model.py`: Model config with async_chunk field

docs/design/feature/hsdp.md

@@ -0,0 +1,146 @@
+# HSDP
+
+This section describes how to add HSDP (Hybrid Sharded Data Parallel) support to a diffusion transformer model. We use the Wan2.2 transformer as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Testing](#testing)
+- [Reference Implementations](#reference-implementations)
+
+---
+
+## Overview
+
+### What is HSDP?
+
+HSDP (Hybrid Sharded Data Parallel) is a memory optimization technique that **shards model weights** across multiple GPUs using PyTorch's FSDP2. Unlike Tensor Parallelism which splits computation, HSDP:
+
+- Shards weights across GPUs to reduce per-GPU memory usage
+- Gathers weights on-demand during forward passes
+- Can work standalone or combined with other parallelism (e.g., Sequence Parallel)
+
+This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
+
+**Important constraints:**
+- HSDP cannot be used with Tensor Parallelism
+- For standalone HSDP (no other parallelism), `hsdp_shard_size` must be specified explicitly
+
+### Architecture
+
+HSDP implementation relies on:
+
+1. **`_hsdp_shard_conditions`**: Model attribute specifying which modules to shard
+2. **`apply_hsdp_to_model`**: Function that applies FSDP2 sharding
+3. **`HSDPInferenceConfig`**: Runtime configuration for HSDP
+
+---
+
+## Step-by-Step Implementation
+
+### Step 1: Identify Modules to Shard
+
+Determine which modules in your transformer should be sharded. Typically, these are:
+
+- Transformer blocks (e.g., `blocks.0`, `blocks.1`, ...)
+- Large submodules with significant weight memory
+
+**Key questions:**
+- Which modules have the largest weights?
+- Which modules are repeated (like transformer blocks)?
+
+### Step 2: Define Shard Conditions
+
+Add `_hsdp_shard_conditions` to your model class. This is a list of functions that return `True` for modules that should be sharded.
+
+**Example (Transformer Blocks):**
+
+```python
+class MyTransformerModel(nn.Module):
+
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        """Match transformer blocks for HSDP sharding.
+
+        Args:
+            name: Module name from named_modules() (e.g., "blocks.0", "blocks.0.attn")
+            module: The module instance
+
+        Returns:
+            True if this module should be sharded
+        """
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_transformer_block]
+```
+
+**Multiple Conditions Example:**
+
+```python
+class MyModel(nn.Module):
+
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    @staticmethod
+    def _is_moe_expert(name: str, module) -> bool:
+        # Also shard MoE expert layers
+        return "experts" in name and name.split(".")[-1].isdigit()
+
+    # Module is sharded if ANY condition returns True
+    _hsdp_shard_conditions = [_is_transformer_block, _is_moe_expert]
+```
+
+---
+
+## Testing
+
+After adding HSDP support, test with:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+parallel_config = DiffusionParallelConfig(
+    use_hsdp=True,
+    hsdp_shard_size=8,  # Shard across 8 GPUs
+)
+omni = Omni(model="your-model-name", parallel_config=parallel_config)
+
+output = omni.generate(
+    "a cup of coffee on the table",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+**Or via command line:**
+
+```bash
+vllm serve Your-org/your-model --omni --port 8091 --use-hsdp
+```
+
+**Verify:**
+
+1. Check logs for "HSDP Inference: replicate_size=..., shard_size=..."
+2. Check logs for "Sharded N modules + root"
+3. Verify memory usage is reduced proportionally
+4. Compare generated output quality with HSDP disabled
+
+---
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Notes |
+|-------|------|-------|
+| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Reference implementation |
+| **HSDP Core** | `vllm_omni/diffusion/distributed/hsdp.py` | `apply_hsdp_to_model`, `shard_model` |
+| **HSDP Tests** | `tests/diffusion/distributed/test_hsdp.py` | Unit tests |
+
+---

docs/getting_started/installation/gpu/cuda.inc.md

@@ -20,20 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-# vllm 0.16.0 is still under prerelease
-uv pip install --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
-
-# vllm 0.16.0 may have some bugs for cuda 12.9, here is how we solve them:
-export FLASHINFER_CUDA_TAG="$(python3 -c 'import torch; print((torch.version.cuda or "12.4").replace(".", ""))')"
-
-uv pip install --upgrade --force-reinstall \
-  "flashinfer-python==0.6.3" \
-  "flashinfer-cubin==0.6.3" \
-  "flashinfer-jit-cache==0.6.3" \
-  --extra-index-url "https://flashinfer.ai/whl/cu${FLASHINFER_CUDA_TAG}"
-
-uv pip install --upgrade --force-reinstall "nvidia-cublas-cu12==12.9.1.4"
-uv pip install --upgrade --force-reinstall "numpy==2.2.6"
+uv pip install vllm --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni

docs/getting_started/installation/npu/npu.inc.md

@@ -33,13 +33,17 @@ docker run --rm \
     -p 8000:8000 \
     -it $IMAGE bash
 
+# Because vllm-ascend will release v0.16.0rc1 after vllm-omni 0.16.0,
+# we have to pin vllm-ascend at the current commit.
+cd /vllm-workspace/vllm-ascend
+git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+pip install -v -e .
+
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.14.0 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.16.0 https://github.com/vllm-project/vllm-omni.git
 
 cd vllm-omni
-sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
- && sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
 pip install -v -e .
 export VLLM_WORKER_MULTIPROC_METHOD=spawn
 ```
@@ -48,13 +52,13 @@ The default workdir is `/workspace`, with vLLM, vLLM-Ascend and vLLM-Omni code p
 
 For other installation methods (pip installation, building from source, custom Docker builds), please refer to the [vllm-ascend installation guide](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
 
-We are keeping [issue #997](https://github.com/vllm-project/vllm-omni/issues/997) up to date with the aligned versions of vLLM, vLLM-Ascend, and vLLM-Omni, and also outlining the Q1 roadmap there.
+We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) up to date with the aligned versions of vLLM, vLLM-Ascend, and vLLM-Omni, and also outlining the Q1 roadmap there.
 
 # --8<-- [end:installation-release]
 
 # --8<-- [start:installation-main]
 
-You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #997](https://github.com/vllm-project/vllm-omni/issues/997) for the status of the latest commit of vLLM-Omni main branch on NPU.)
+You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
 
 ```bash
 # Pin vLLM version to 0.16.0
@@ -62,12 +66,14 @@ cd /vllm-workspace/vllm
 git pull origin main
 git fetch origin --tags
 git checkout v0.16.0
+VLLM_TARGET_DEVICE=empty pip install -v -e .
 
 # Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time.
 cd ../vllm-ascend
 git pull origin main
 git fetch origin --tags
 git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+pip install -v -e .
 
 # Install vLLM-Omni from the latest main branch
 cd ../vllm-omni
@@ -94,10 +100,10 @@ Here's an example deployment command that has been verified on 4 x NPUs:
 
 ```bash
 # Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-omni:v0.14.0
+# export IMAGE=quay.io/ascend/vllm-omni:v0.16.0
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-omni:v0.14.0-a3
-export IMAGE=quay.io/ascend/vllm-omni:v0.14.0
+# export IMAGE=quay.io/ascend/vllm-omni:v0.16.0-a3
+export IMAGE=quay.io/ascend/vllm-omni:v0.16.0
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \

docs/getting_started/quickstart.md

@@ -19,18 +19,7 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-# vllm 0.16.0 is still under prerelease
-uv pip install --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
-# vllm 0.16.0 may have some bugs for cuda 12.9, here is how we solve them:
-export FLASHINFER_CUDA_TAG="$(python3 -c 'import torch; print((torch.version.cuda or "12.4").replace(".", ""))')"
-uv pip install --upgrade --force-reinstall \
-  "flashinfer-python==0.6.3" \
-  "flashinfer-cubin==0.6.3" \
-  "flashinfer-jit-cache==0.6.3" \
-  --extra-index-url "https://flashinfer.ai/whl/cu${FLASHINFER_CUDA_TAG}"
-uv pip install --upgrade --force-reinstall "nvidia-cublas-cu12==12.9.1.4"
-uv pip install --upgrade --force-reinstall "numpy==2.2.6"
-
+uv pip install vllm==0.16.0 --torch-backend=auto
 
 # On ROCm
 uv pip install vllm==0.16.0 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700

docs/mkdocs/hooks/generate_argparse.py

@@ -124,6 +124,7 @@ def extract_omni_serve_subparser_init():
                         "logger": logger,
                         "DummySubparsers": DummySubparsers,
                         "argparse": __import__("argparse"),
+                        "json": __import__("json"),
                         "DESCRIPTION": DESCRIPTION,
                     }
                     exec(code, exec_globals, local_vars)

docs/models/supported_models.md

@@ -41,7 +41,9 @@ th {
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |
+|`GlmImageForConditionalGeneration` | GLM-Image | `zai-org/GLM-Image` |
 |`NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` |
+|`MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` |
 
 
 ## List of Supported Models for NPU
@@ -64,7 +66,10 @@ th {
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` |
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
+| `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
+| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` |
 |`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
+|`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
 |`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -16,6 +16,8 @@ The following parallelism methods are currently supported in vLLM-Omni:
 
 5. [VAE Patch Parallelism](#vae-patch-parallelism): VAE patch parallelism shards VAE decode spatially across ranks. This can reduce the peak memory of VAE decode and (depending on resolution and communication overhead) speed up VAE decode.
 
+6. [HSDP](#hsdp): Hybrid Sharded Data Parallel shards model weights across GPUs using PyTorch FSDP2. This reduces per-GPU memory usage, enabling inference of large models on GPUs with limited memory.
+
 The following table shows which models are currently supported by parallelism method:
 
 ### ImageGen
@@ -30,7 +32,7 @@ The following table shows which models are currently supported by parallelism me
 | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
 | **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
 | **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |         ✅          |
-| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |
+| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ✅        |         ✅          |
 | **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |
 | **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ✅       |        ✅        |         ❌          |
 
@@ -49,9 +51,9 @@ The following table shows which models are currently supported by parallelism me
 
 ### VideoGen
 
-| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel |
-|-------|------------------|------------|---------|--------------------------|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ |
+| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP |
+|-------|------------------|:----------:|:--------------:|:---------------:|:----:|
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ |
 
 ### Tensor Parallelism
 
@@ -83,7 +85,7 @@ outputs = omni.generate(
 VAE patch parallelism distributes the VAE decode workload across multiple ranks by splitting the latent spatially. It is configured via `DiffusionParallelConfig.vae_patch_parallel_size` and can be combined with other parallelism methods (e.g., TP).
 
 !!! note "Enablement and feature gate"
-    - VAE patch parallelism is currently **enabled only for validated pipelines** (currently: `Tongyi-MAI/Z-Image-Turbo`).
+    - VAE patch parallelism is currently **enabled only for validated pipelines** (currently: `Tongyi-MAI/Z-Image-Turbo`, `stabilityai/stable-diffusion-3.5`).
     - If `vae_patch_parallel_size > 1` is set for a validated pipeline, vLLM-Omni will automatically enable `vae_use_tiling` as a safety gate. (We use `vae_use_tiling` because it indicates the VAE supports diffusers tiling parameters like `tile_latent_min_size` and `tile_overlap_factor`.)
 
 #### Offline Inference
@@ -330,3 +332,84 @@ You can enable CFG-Parallel in online serving for diffusion models via `--cfg-pa
 ```bash
 vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2
 ```
+
+### HSDP
+
+HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
+
+Unlike Tensor Parallelism which splits computation, HSDP uses PyTorch's FSDP2 to shard and redistribute weights at runtime. Each GPU only holds a fraction of the model weights, and weights are gathered on-demand during forward passes.
+
+#### Configuration
+
+HSDP is configured via `DiffusionParallelConfig`:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `use_hsdp` | bool | False | Enable HSDP |
+| `hsdp_shard_size` | int | -1 | Number of GPUs to shard weights across. -1 = auto (requires other parallelism > 1) |
+| `hsdp_replicate_size` | int | 1 | Number of replica groups. Each group holds a full sharded copy |
+
+**Constraints:**
+
+- `hsdp_replicate_size × hsdp_shard_size == world_size`
+- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
+
+#### Operating Modes
+
+HSDP can work in two modes:
+
+- **Standalone Mode**: HSDP alone without other parallelism. Must specify `hsdp_shard_size` explicitly.
+- **Combined Mode**: HSDP overlays on top of other parallelism (Ulysses Sequence Parallel, CFG Parallel). HSDP dimensions must match world_size.
+
+#### Offline Inference
+
+**Standalone HSDP** (shard across 4 GPUs, no other parallelism):
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    parallel_config=DiffusionParallelConfig(
+        use_hsdp=True,
+        hsdp_shard_size=4,  # Shard across 4 GPUs
+    ),
+)
+
+outputs = omni.generate(
+    "A cat playing piano",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+**Combined HSDP + Sequence Parallel**:
+
+```python
+omni = Omni(
+    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    parallel_config=DiffusionParallelConfig(
+        ulysses_degree=4,  # Sequence parallel
+        use_hsdp=True,     # HSDP overlays on SP
+    ),
+)
+```
+
+#### Online Serving
+
+**Standalone HSDP** (shard model across 4 GPUs):
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --hsdp-shard-size 4
+```
+
+**Combined with Sequence Parallel**:
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --usp 4
+```
+
+#### Adding HSDP Support to New Models
+
+For detailed instructions on adding HSDP support to new models, see the [HSDP Contributing Guide](../../design/feature/hsdp.md).

docs/user_guide/diffusion/quantization/gguf.md

@@ -0,0 +1,185 @@
+# GGUF Quantization
+
+## Goals
+1. Reuse vLLM quantization configs and weight loaders as much as possible.
+2. Add native GGUF support to diffusion transformers without changing model definitions.
+3. Keep user-facing knobs minimal and consistent across offline and online flows.
+
+## Scope
+1. Models: Z-Image, and Flux2-klein.
+2. Components: diffusion transformer weights, loader paths, and quantization configs.
+3. Modes: native GGUF (transformer-only weights).
+
+## Architecture Overview
+1. `OmniDiffusionConfig` accepts `quantization` or `quantization_config`.
+2. Diffusion quantization wrapper (`DiffusionGgufConfig`) produces vLLM `QuantizationConfig` objects for linear layers.
+3. `DiffusersPipelineLoader` branches on quantization method and loads either HF weights or GGUF weights for the transformer.
+4. GGUF transformer loading is routed through model-specific adapters (e.g., Flux2Klein).
+5. vLLM GGUF path uses `GGUFConfig` and `GGUFLinearMethod` for matmul.
+
+## Call Chain (Offline)
+```
+CLI (examples/offline_inference/text_to_image/text_to_image.py)
+  |
+  v
+Omni (vllm_omni/entrypoints/omni.py)
+  |
+  v
+OmniStage (diffusion)
+  |
+  v
+DiffusionWorker
+  |
+  v
+DiffusionModelRunner
+  |
+  v
+DiffusersPipelineLoader
+  |
+  v
+Pipeline.forward (Flux2/Qwen/Z-Image)
+  |
+  v
+DiffusionEngine
+  |
+  v
+OmniRequestOutput
+  |
+  v
+Client (saved PNG)
+```
+
+## Call Chain (Online)
+```
+Client
+  |
+  | POST /v1/images/generations
+  v
+APIServer (vllm_omni/entrypoints/openai/api_server.py)
+  |
+  v
+_generate_with_async_omni
+  |
+  v
+AsyncOmni
+  |
+  v
+DiffusionEngine
+  |
+  v
+OmniRequestOutput
+  |
+  v
+encode_image_base64
+  |
+  v
+ImageGenerationResponse
+  |
+  v
+Client
+```
+
+## Call Chain (GGUF Operator Path)
+```
+Pipeline.forward (Flux2/Qwen/Z-Image)
+  |
+  v
+Transformer blocks
+  |
+  v
+QKVParallelLinear / ColumnParallelLinear / RowParallelLinear
+  |
+  v
+LinearBase.forward
+  |
+  v
+QuantMethod.apply (GGUFLinearMethod.apply)
+  |
+  v
+fused_mul_mat_gguf
+  |
+  v
+_fused_mul_mat_gguf (custom op)
+  |
+  v
+ops.ggml_dequantize
+  |
+  v
+x @ weight.T
+```
+
+## GGUF Weight Loading Path (Transformer-Only)
+1. `DiffusersPipelineLoader.load_model` detects `quantization_config.method == "gguf"`.
+2. `gguf_model` is resolved as one of: local file, `repo/file.gguf`, or `repo:quant_type`.
+3. GGUF weights are routed through adapters in `vllm_omni/diffusion/model_loader/gguf_adapters/`.
+4. Name mapping is applied per-architecture (Z-Image, Flux2Klein).
+5. GGUF weights are loaded into transformer modules, remaining non-transformer weights come from the HF checkpoint.
+
+## GGUF Adapter Design
+1. `GGUFAdapter` is an abstract base class for model-specific adapters.
+2. `Flux2KleinGGUFAdapter` implements Flux2-Klein remapping + qkv split + adaLN swap.
+3. `ZImageGGUFAdapter` implements Z-Image qkv + ffn shard handling and linear qweight routing.
+4. `get_gguf_adapter(...)` strictly selects by model class/config; unsupported models raise an error (no fallback adapter).
+
+Adapter paths:
+- Base: `vllm_omni/diffusion/model_loader/gguf_adapters/base.py`
+- Z-Image: `vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py`
+- Flux2-Klein: `vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py`
+
+## User Usage (Offline)
+
+### Baseline BF16
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model /workspace/models/black-forest-labs/FLUX.2-klein-4B \
+  --prompt "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture" \
+  --height 768 \
+  --width 1360 \
+  --seed 42 \
+  --cfg_scale 4.0 \
+  --num_images_per_prompt 1 \
+  --num_inference_steps 4 \
+  --output outputs/flux2_klein_4b.png
+```
+
+### Native GGUF (Transformer Only)
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model /workspace/models/black-forest-labs/FLUX.2-klein-4B \
+  --gguf-model "/workspace/models/unsloth/FLUX.2-klein-4B-GGUF/flux-2-klein-4b-Q8_0.gguf" \
+  --quantization gguf \
+  --prompt "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture" \
+  --height 768 \
+  --width 1360 \
+  --seed 42 \
+  --cfg_scale 4.0 \
+  --num_images_per_prompt 1 \
+  --num_inference_steps 4 \
+  --output outputs/flux2_klein_4b_gguf.png
+```
+
+Notes for GGUF:
+1. Many GGUF repos do not ship `model_index.json` and configs. Use the base repo for `--model` and only pass the GGUF file via `--gguf-model`.
+2. `gguf_model` supports local path, `repo/file.gguf`, or `repo:quant_type`.
+
+## User Usage (Online)
+
+### Start Server (Native GGUF via CLI)
+```bash
+vllm serve /workspace/models/black-forest-labs/FLUX.2-klein-4B \
+  --omni \
+  --port 8000 \
+  --quantization-config '{"method":"gguf","gguf_model":"/workspace/models/unsloth/FLUX.2-klein-4B-GGUF/flux-2-klein-4b-Q8_0.gguf"}'
+```
+
+### Online Request (Images API)
+```bash
+curl -X POST http://localhost:8000/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "prompt": "a dragon laying over the spine of the Green Mountains of Vermont",
+    "size": "1024x1024",
+    "seed": 42,
+    "num_inference_steps": 4
+  }'
+```

docs/user_guide/diffusion/quantization/overview.md

@@ -7,6 +7,7 @@ vLLM-Omni supports quantization of DiT linear layers to reduce memory usage and
 | Method | Guide |
 |--------|-------|
 | FP8 | [FP8](fp8.md) |
+| GGUF | [GGUF](gguf.md) |
 
 ## Device Compatibility
 

docs/user_guide/diffusion_acceleration.md

@@ -30,6 +30,8 @@ vLLM-Omni also supports parallelism methods for diffusion models, including:
 
 5. [VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism) - shards VAE decode/encode spatially across ranks to reduce VAE peak memory (and can speed up VAE decode).
 
+6. [HSDP](diffusion/parallelism_acceleration.md#hsdp) - Hybrid Sharded Data Parallel shards model weights across GPUs to reduce per-GPU memory usage, enabling inference of large models on limited GPU memory.
+
 ## Quick Comparison
 
 ### Cache Methods
@@ -69,9 +71,9 @@ The following table shows which models are currently supported by each accelerat
 
 ### VideoGen
 
-| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel |
-|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:----------------:|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ |
+| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP |
+|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ |
 
 ### Quantization
 
@@ -265,6 +267,30 @@ outputs = omni.generate(
 )
 ```
 
+### Using HSDP
+
+HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+    parallel_config=DiffusionParallelConfig(
+        use_hsdp=True,
+        hsdp_replicate_size=1,  # Number of replica groups
+        hsdp_shard_size=8,      # Number of GPUs to shard across (or -1 for auto)
+    ),
+)
+
+outputs = omni.generate(
+    "A cat playing piano",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
 ### Using CFG-Parallel
 
 Run image-to-image:
@@ -319,3 +345,4 @@ For detailed information on each acceleration method:
 - **[Sequence Parallelism](diffusion/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration.
 - **[CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel)** - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks.
 - **[VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism)** - Guidance on how to reduce VAE memory via patch/tile parallelism.
+- **[HSDP](diffusion/parallelism_acceleration.md#hsdp)** - Hybrid Sharded Data Parallel for memory-efficient inference of large models.

docs/user_guide/examples/offline_inference/bagel.md

@@ -154,6 +154,24 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
 
 ------
 
+#### Tensor Parallelism (TP)
+
+For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) by modifying the stage configuration (e.g., [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)).
+
+1. **Set `tensor_parallel_size`**: Increase this value (e.g., to `2` or `4`).
+2. **Set `devices`**: Specify the comma-separated GPU IDs to be used for the stage (e.g., `"0,1"`).
+
+Example configuration for TP=2 on GPUs 0 and 1:
+```yaml
+    engine_args:
+      tensor_parallel_size: 2
+      ...
+    runtime:
+      devices: "0,1"
+```
+
+------
+
 #### 🔗 Runtime Configuration
 
 | Parameter             | Value   | Description                      |
@@ -162,6 +180,74 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
 | `max_inflight`        | `1`     | Maximum inflight requests        |
 | `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB)   |
 
+## Using Mooncake Connector
+
+[Mooncake](https://github.com/kvcache-ai/Mooncake) is a high-performance distributed KV cache transfer engine that enables efficient cross-node data movement via TCP or RDMA, making it ideal for multi-node disaggregated inference.
+
+By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can switch to the Mooncake connector for better performance on multi-GPU setups and to enable multi-node deployment.
+
+### Prerequisites
+
+Install the Mooncake transfer engine:
+
+```bash
+# For CUDA-enabled systems (recommended)
+pip install mooncake-transfer-engine
+
+# For non-CUDA systems
+pip install mooncake-transfer-engine-non-cuda
+```
+
+### Step 1: Start the Mooncake Master
+
+On the **primary node**, start the Mooncake master service (run in a separate terminal or background with `&`):
+
+```bash
+# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
+# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
+mkdir -p ./mc_storage
+
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=0.0.0.0 \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003 \
+  --root_fs_dir=./mc_storage/ \
+  --cluster_id=mc-local-1 &
+```
+
+### Step 2: Run Offline Inference with Mooncake
+
+Use the provided Mooncake stage config [`bagel_multiconnector.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml). Before launching, update the `metadata_server` and `master` addresses in the YAML to match your Mooncake master node's IP (use `127.0.0.1` for single-node testing).
+
+```bash
+cd examples/offline_inference/bagel
+
+# Text to Image with Mooncake
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2img \
+                  --prompts "A cute cat" \
+                  --stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+
+# Image to Text with Mooncake
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality img2text \
+                  --image-path /path/to/image.jpg \
+                  --prompts "Describe this image" \
+                  --stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+
+# Text to Text with Mooncake
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2text \
+                  --prompts "What is the capital of France?" \
+                  --stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+```
+
+For more details on the Mooncake connector and multi-node setup, see the [Mooncake Store Connector documentation](https://github.com/vllm-project/vllm-omni/tree/main/docs/design/feature/omni_connectors/mooncake_store_connector.md).
+
+------
+
 ## FAQ
 
 - If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.

docs/user_guide/examples/offline_inference/glm_image.md

@@ -0,0 +1,156 @@
+# GLM-Image Multistage End-to-End Inference
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/glm_image>.
+
+
+This example demonstrates how to run GLM-Image with the vLLM-Omni multistage architecture.
+
+## Architecture
+
+GLM-Image uses a 2-stage pipeline:
+
+```
+┌─────────────────────────────────────────────────────────────┐
+│                     GLM-Image Pipeline                       │
+├─────────────────────────────────────────────────────────────┤
+│                                                             │
+│  Stage 0 (AR Model)              Stage 1 (Diffusion)        │
+│  ┌─────────────────┐            ┌─────────────────────┐     │
+│  │ vLLM-optimized  │            │  GlmImagePipeline   │     │
+│  │ GlmImageFor     │  prior     │  ┌───────────────┐  │     │
+│  │ Conditional     │──tokens───►│  │ DiT Denoiser  │  │     │
+│  │ Generation      │            │  └───────────────┘  │     │
+│  │ (9B AR model)   │            │         │          │     │
+│  └─────────────────┘            │         ▼          │     │
+│         ▲                       │  ┌───────────────┐  │     │
+│         │                       │  │  VAE Decode   │──┼──► Image
+│    Text/Image                   │  └───────────────┘  │     │
+│      Input                      └─────────────────────┘     │
+│                                                             │
+└─────────────────────────────────────────────────────────────┘
+```
+
+## Features
+
+- **vLLM-optimized AR**: Uses PagedAttention and tensor parallelism for faster prior token generation
+- **Flexible deployment**: AR and Diffusion stages can run on different GPUs
+- **Text-to-Image**: Generate images from text descriptions
+- **Image-to-Image**: Edit existing images with text prompts
+
+## Usage
+
+### Text-to-Image
+
+```bash
+python end2end.py \
+    --model-path /path/to/glm-image \
+    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --prompt "A beautiful sunset over the ocean with sailing boats" \
+    --height 1024 \
+    --width 1024 \
+    --output output_t2i.png
+```
+
+### Image-to-Image (Image Editing)
+
+```bash
+python end2end.py \
+    --model-path /path/to/glm-image \
+    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --prompt "Transform this scene into a winter wonderland" \
+    --image input.png \
+    --output output_i2i.png
+```
+
+### With Custom Parameters
+
+```bash
+python end2end.py \
+    --model-path /path/to/glm-image \
+    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --prompt "A photorealistic cat sitting on a window sill" \
+    --height 1024 \
+    --width 1024 \
+    --num-inference-steps 50 \
+    --guidance-scale 1.5 \
+    --seed 42 \
+    --output output.png
+```
+
+## Shell Scripts
+
+### Run Text-to-Image
+
+```bash
+./run_t2i.sh
+```
+
+### Run Image-to-Image
+
+```bash
+./run_i2i.sh --image /path/to/input.png
+```
+
+## Stage Configuration
+
+The stage config (`glm_image.yaml`) defines:
+
+- **Stage 0 (AR)**: Uses `GPUARWorker` with vLLM engine
+
+  - Model: `GlmImageForConditionalGeneration`
+  - Output: `token_ids` (prior tokens)
+
+- **Stage 1 (Diffusion)**: Uses diffusion engine
+  - Model: `GlmImagePipeline`
+  - Output: Generated image
+
+See `vllm_omni/model_executor/stage_configs/glm_image.yaml` for full configuration.
+
+## Comparison with Single-Stage
+
+| Aspect      | Single-Stage (transformers) | Multistage (vLLM)   |
+| ----------- | --------------------------- | ------------------- |
+| AR Model    | transformers native         | vLLM PagedAttention |
+| Memory      | Higher (no KV cache opt)    | Lower (optimized)   |
+| Throughput  | Lower                       | Higher              |
+| Flexibility | Single GPU                  | Multi-GPU support   |
+
+## Troubleshooting
+
+### OOM Error
+
+Try reducing memory usage:
+
+```bash
+# In glm_image.yaml, adjust:
+gpu_memory_utilization: 0.5  # Reduce from 0.6
+```
+
+### Slow Initialization
+
+The first run loads model weights. Subsequent runs are faster:
+
+```bash
+--stage-init-timeout 900  # Increase timeout for slow storage
+```
+
+## Requirements
+
+- vLLM-Omni with GLM-Image support
+- CUDA-capable GPU (recommended: H100/A100 with 80GB)
+- GLM-Image model weights
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/glm_image/end2end.py"
+    ``````
+??? abstract "run_i2i.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/glm_image/run_i2i.sh"
+    ``````
+??? abstract "run_t2i.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/glm_image/run_t2i.sh"
+    ``````

docs/user_guide/examples/offline_inference/image_to_video.md

@@ -69,6 +69,11 @@ Key arguments:
 - `--vae-use-tiling`: Enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+- `--use-hsdp`: Enable Hybrid Sharded Data Parallel to shard model weights across GPUs.
+- `--hsdp-shard-size`: Number of GPUs to shard model weights across within each replica group. -1 (default) auto-calculates as world_size / replicate_size.
+- `--hsdp-replicate-size`: Number of replica groups for HSDP. Each replica holds a full sharded copy. Default 1 means pure sharding (no replication).
+
+
 
 > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
 

docs/user_guide/examples/offline_inference/qwen3_tts.md

@@ -90,13 +90,43 @@ Examples:
 python end2end.py --query-type Base --mode-tag icl
 ```
 
+## Streaming Mode
+
+Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config):
+
+```bash
+python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_stream
+```
+
+Each 25-frame Code2Wav chunk is logged as it arrives. The final WAV file is written once generation
+completes. This demonstrates that audio data is available progressively rather than only at the end.
+
+> **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged.
+
+## Batched Decoding
+
+The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_batch_size > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
+
+```
+python end2end.py --query-type CustomVoice \
+    --txt-prompts benchmark_prompts.txt \
+    --batch-size 4 \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml
+```
+
+**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_batch_size >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_batch_size`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
+
 ## Notes
 
 - The script uses the model paths embedded in `end2end.py`. Update them if your local cache path differs.
-- Use `--output-dir` (preferred) or `--output-wav` to change the output folder.
+- Use `--output-dir` to change the output folder.
 
 ## Example materials
 
+??? abstract "benchmark_prompts.txt"
+    ``````txt
+    --8<-- "examples/offline_inference/qwen3_tts/benchmark_prompts.txt"
+    ``````
 ??? abstract "end2end.py"
     ``````py
     --8<-- "examples/offline_inference/qwen3_tts/end2end.py"

docs/user_guide/examples/offline_inference/text_to_image.md

@@ -3,7 +3,7 @@
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_image>.
 
 
-This folder provides several entrypoints for experimenting with `Qwen/Qwen-Image` `Qwen/Qwen-Image-2512` `Tongyi-MAI/Z-Image-Turbo` using vLLM-Omni:
+This folder provides several entrypoints for experimenting with `Qwen/Qwen-Image` `Qwen/Qwen-Image-2512` `Tongyi-MAI/Z-Image-Turbo` `stepfun-ai/NextStep-1.1` using vLLM-Omni, note that NextStep-1.1 has different architecture so we treat it differently regarding running arguments and pipeline.
 
 - `text_to_image.py`: command-line script for single image generation with advanced options.
 - `web_demo.py`: lightweight Gradio UI for interactive prompt/seed/CFG exploration.
@@ -77,6 +77,8 @@ if __name__ == "__main__":
 
 ## Local CLI Usage
 
+### Qwen/Tongyi Models
+
 ```bash
 python text_to_image.py \
   --model Tongyi-MAI/Z-Image-Turbo \
@@ -90,7 +92,26 @@ python text_to_image.py \
   --output outputs/coffee.png
 ```
 
-Key arguments:
+### NextStep Models
+
+NextStep-1.1 can have extra arguments
+```bash
+python text_to_image.py \
+  --model stepfun-ai/NextStep-1.1 \
+  --prompt "A baby panda wearing an Iron Man mask, holding a board with 'NextStep-1' written on it" \
+  --height 512 \
+  --width 512 \
+  --num-inference-steps 28 \
+  --guidance-scale 7.5 \
+  --guidance-scale-2 1.0 \
+  --cfg-schedule constant \
+  --output nextstep_output.png \
+  --seed 42
+```
+
+### Key Arguments
+
+**Common arguments:**
 
 - `--prompt`: text description (string).
 - `--seed`: integer seed for deterministic sampling.
@@ -101,8 +122,15 @@ Key arguments:
 - `--output`: path to save the generated PNG.
 - `--vae-use-slicing`: enable VAE slicing for memory optimization.
 - `--vae-use-tiling`: enable VAE tiling for memory optimization.
-- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion_acceleration.md#using-cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+- `--guidance-scale`: classifier-free guidance scale.
+
+**NextStep-1.1 specific:**
+- `--guidance-scale-2`: secondary guidance scale, e.g. image-level CFG (default: 1.0).
+- `--timesteps-shift`: timesteps shift parameter for sampling (default: 1.0).
+- `--cfg-schedule`: CFG schedule type, "constant" or "linear" (default: "constant").
+- `--use-norm`: apply layer normalization to sampled tokens.
 
 > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
 

docs/user_guide/examples/online_serving/bagel.md

@@ -25,12 +25,152 @@ cd /workspace/vllm-omni/examples/online_serving/bagel
 bash run_server.sh
 ```
 
-If you have a custom stage configs file, launch the server with the command below:
-
 ```bash
 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```
 
+#### 🚀 Tensor Parallelism (TP)
+
+For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) for the server.
+
+1. **Modify Stage Config**: Create or modify a stage configuration yaml (e.g., [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)). Set `tensor_parallel_size` to `2` (or more) and update `devices` to include multiple GPU IDs (e.g., `"0,1"`).
+
+```yaml
+    engine_args:
+      tensor_parallel_size: 2
+      ...
+    runtime:
+      devices: "0,1"
+```
+
+2. **Launch Server**:
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/your/custom_bagel.yaml
+```
+
+#### Using Mooncake Connector
+
+By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can use the [Mooncake](https://github.com/kvcache-ai/Mooncake) connector to transfer KV cache between stages, which also enables multi-node deployment.
+
+**1. Install Mooncake**
+
+```bash
+# For CUDA-enabled systems (recommended)
+pip install mooncake-transfer-engine
+
+# For non-CUDA systems
+pip install mooncake-transfer-engine-non-cuda
+```
+
+**2. Start Mooncake Master** on the primary node:
+
+```bash
+# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
+# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
+mkdir -p ./mc_storage
+
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=0.0.0.0 \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003 \
+  --root_fs_dir=./mc_storage/ \
+  --cluster_id=mc-local-1 &
+```
+
+**3. Launch the server** with the Mooncake stage config:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+```
+
+> **Note**: Before launching, edit [`bagel_multiconnector.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml) and replace the `metadata_server` and `master` addresses with your Mooncake master node's actual IP. For single-node testing, `127.0.0.1` works.
+
+The client-side usage is identical to the default setup -- the Mooncake connector is transparent to the API. See the requests section below.
+
+For more details on the Mooncake connector configuration, see the [Mooncake Store Connector documentation](https://github.com/vllm-project/vllm-omni/tree/main/docs/design/feature/omni_connectors/mooncake_store_connector.md).
+
+#### Multi-Node Deployment
+
+You can deploy each stage on a **separate node** for better resource utilization. In this example, the orchestrator (Stage 0 / Thinker) and Stage 1 (DiT) run on different machines, connected via Mooncake.
+
+Replace `<ORCHESTRATOR_IP>` below with the actual IP address of your orchestrator node (e.g., `10.244.227.244`).
+
+> [!WARNING]
+> **Before launching**, edit [`bagel_multiconnector.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml) and replace the `metadata_server` and `master` addresses with your Mooncake master node's actual IP. Mismatched addresses will cause silent connection failures.
+
+**1. Start Mooncake Master** (on the orchestrator node):
+
+```bash
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=<ORCHESTRATOR_IP> \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003
+```
+
+**2. Launch Stage 0 (Thinker / Orchestrator)** on the orchestrator node:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
+    --port 8000 \ # API server port for client requests
+    --stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml \
+    --stage-id 0 \
+    -oma <ORCHESTRATOR_IP> \
+    -omp 8091
+```
+
+**3. Launch Stage 1 (DiT)** on the remote node in headless mode:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml \
+    --stage-id 1 \
+    --headless \
+    -oma <ORCHESTRATOR_IP> \
+    -omp 8091
+```
+
+**Mooncake Master arguments:**
+
+| Argument | Description |
+| :------- | :---------- |
+| `--rpc_port` | Mooncake RPC port for control-plane coordination between stages |
+| `--enable_http_metadata_server` | Enable the HTTP metadata server for service discovery |
+| `--http_metadata_server_host` | IP address to bind the metadata server (use the orchestrator node's IP) |
+| `--http_metadata_server_port` | Port for the HTTP metadata server |
+| `--metrics_port` | Port for Prometheus-compatible metrics endpoint |
+
+**vllm serve arguments:**
+
+| Argument | Description |
+| :------- | :---------- |
+| `--stage-id` | Which stage this process runs (0 = Thinker, 1 = DiT) |
+| `--headless` | Run without the API server (worker-only mode) |
+| `-oma` | Orchestrator master address |
+| `-omp` | Orchestrator master port for Stage 1 to connect to Stage 0 for task coordination |
+
+> [!IMPORTANT]
+> **Startup Order**: Stage 0 (orchestrator) must be launched **before** Stage 1 (headless).
+> Stage 0 will appear to hang on startup until Stage 1 (worker) connects — this is expected behavior.
+
+**Network Requirements**
+
+All nodes must have network connectivity to each other. Ensure the following ports are open **between all participating nodes**:
+
+| Port | Protocol | Service | Direction |
+| :--- | :------- | :------ | :-------- |
+| 50051 | TCP | Mooncake Master RPC | Worker → Orchestrator |
+| 8080 | TCP | Mooncake HTTP Metadata Server | Worker → Orchestrator |
+| 8091 | TCP | Orchestrator Master (`-omp`) | Worker → Orchestrator |
+| 8000 | TCP | API Server (`--port`) | Client → Orchestrator |
+| 9003 | TCP | Metrics (optional) | Monitoring → Orchestrator |
+
+> **Tip**: If nodes are behind a firewall or in different VPCs/security groups, make sure the above ports are allowed in ingress/egress rules. All nodes should be reachable via their IP addresses (no NAT). Using nodes on the same subnet or VPC is recommended to minimize latency for Mooncake KV cache transfers.
+
 ### Send Multi-modal Request
 
 Get into the bagel folder:
@@ -242,3 +382,7 @@ sudo apt install ffmpeg
     ``````sh
     --8<-- "examples/online_serving/bagel/run_server.sh"
     ``````
+??? abstract "run_server_stage_cli.sh"
+    ``````sh
+    --8<-- "examples/online_serving/bagel/run_server_stage_cli.sh"
+    ``````

docs/user_guide/examples/online_serving/glm_image.md

@@ -0,0 +1,269 @@
+# GLM-Image Online Serving
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/glm_image>.
+
+
+This example demonstrates how to deploy GLM-Image for online image generation using vLLM-Omni.
+
+## 🛠️ Installation
+
+Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
+
+## Run examples (GLM-Image)
+
+**Note**: These examples work with the default configuration on **2× NVIDIA A100 (80GB)** or equivalent. Stage 0 (AR) and Stage 1 (Diffusion) each use one GPU by default. For single-GPU setups, modify the stage configuration to share the same device.
+
+### Launch the Server
+
+```bash
+# Use default configuration
+vllm serve zai-org/GLM-Image --omni --port 8091
+```
+
+Or use the convenience script:
+
+```bash
+cd examples/online_serving/glm_image
+bash run_server.sh
+```
+
+If you have a custom stage configs file:
+
+```bash
+vllm serve zai-org/GLM-Image --omni --port 8091 --stage-configs-path /path/to/glm_image.yaml
+```
+
+### Send Requests
+
+Get into the glm_image folder:
+
+```bash
+cd examples/online_serving/glm_image
+```
+
+Send request via Python:
+
+```bash
+python openai_chat_client.py --prompt "A cute cat sitting on a window sill"
+```
+
+The Python client supports the following command-line arguments:
+
+- `--prompt` (or `-p`): Text prompt for generation (default: `A beautiful sunset over the ocean with sailing boats`)
+- `--output` (or `-o`): Output file path (default: `glm_image_output.png`)
+- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
+- `--image` (or `-i`): Input image path (for image-to-image editing)
+- `--height`: Image height in pixels (default: 1024)
+- `--width`: Image width in pixels (default: 1024)
+- `--steps`: Number of inference steps (default: 50)
+- `--guidance-scale`: Classifier-free guidance scale (default: 1.5)
+- `--seed`: Random seed (default: 42)
+- `--negative`: Negative prompt
+
+## Modality Control
+
+GLM-Image supports **text-to-image** and **image-to-image** modes.
+
+The default yaml configuration deploys AR on GPU 0 and DiT on GPU 1. You can use the default configuration file: [`glm_image.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/glm_image.yaml)
+
+| Mode           | Input        | Output | Description                        |
+| -------------- | ------------ | ------ | ---------------------------------- |
+| Text-to-Image  | Text         | Image  | Generate images from text prompts  |
+| Image-to-Image | Image + Text | Image  | Edit images with text instructions |
+
+### Text-to-Image
+
+Generate images from text prompts:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "A photorealistic mountain landscape at sunset" \
+    --height 1024 \
+    --width 1024 \
+    --output landscape.png
+```
+
+**Using curl**
+
+```bash
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [
+      {"role": "user", "content": "A beautiful sunset over the ocean with sailing boats"}
+    ],
+    "extra_body": {
+      "height": 1024,
+      "width": 1024,
+      "num_inference_steps": 50,
+      "guidance_scale": 1.5,
+      "seed": 42
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
+```
+
+Or use the script:
+
+```bash
+bash run_curl_text_to_image.sh "A futuristic city skyline at night"
+```
+
+### Image-to-Image (Image Editing)
+
+Edit images with text instructions:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "Convert this image to watercolor style" \
+    --image input.png \
+    --output watercolor.png
+```
+
+**Using curl**
+
+```bash
+IMG_B64=$(base64 < input.png | tr -d '\n')
+
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d @- <<EOF | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
+{
+  "messages": [{
+    "role": "user",
+    "content": [
+      {"type": "text", "text": "Convert this image to watercolor style"},
+      {"type": "image_url", "image_url": {"url": "data:image/png;base64,'$IMG_B64'"}}
+    ]
+  }],
+  "extra_body": {
+    "height": 1024,
+    "width": 1024,
+    "num_inference_steps": 50,
+    "guidance_scale": 1.5,
+    "seed": 42
+  }
+}
+EOF
+```
+
+Or use the script:
+
+```bash
+bash run_curl_image_edit.sh input.png "Convert to watercolor style"
+```
+
+## Generation Parameters (extra_body)
+
+| Parameter             | Type  | Default | Description                         |
+| --------------------- | ----- | ------- | ----------------------------------- |
+| `height`              | int   | 1024    | Image height in pixels              |
+| `width`               | int   | 1024    | Image width in pixels               |
+| `num_inference_steps` | int   | 50      | Number of diffusion denoising steps |
+| `guidance_scale`      | float | 1.5     | Classifier-free guidance scale      |
+| `seed`                | int   | 42      | Random seed for reproducibility     |
+| `negative_prompt`     | str   | None    | Negative prompt                     |
+
+## Response Format
+
+```json
+{
+  "id": "chatcmpl-xxx",
+  "created": 1234567890,
+  "model": "zai-org/GLM-Image",
+  "choices": [
+    {
+      "index": 0,
+      "message": {
+        "role": "assistant",
+        "content": [
+          {
+            "type": "image_url",
+            "image_url": {
+              "url": "data:image/png;base64,..."
+            }
+          }
+        ]
+      },
+      "finish_reason": "stop"
+    }
+  ],
+  "usage": {}
+}
+```
+
+## Extract Image
+
+```bash
+# From a saved JSON response
+cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
+```
+
+## Architecture
+
+GLM-Image uses a 2-stage multistage pipeline:
+
+```
+Stage 0 (AR Model)                Stage 1 (Diffusion)
+┌───────────────────┐            ┌─────────────────────┐
+│ vLLM-optimized    │  prior     │  GlmImagePipeline   │
+│ GlmImageFor       │──tokens──►│  ┌───────────────┐  │
+│ Conditional       │            │  │ DiT Denoiser  │  │
+│ Generation        │            │  └───────┬───────┘  │
+│ (9B AR model)     │            │          ▼          │
+└───────────────────┘            │  ┌───────────────┐  │
+        ▲                        │  │  VAE Decode   │──┼──► Image
+        │                        │  └───────────────┘  │
+   Text / Image                  └─────────────────────┘
+     Input
+```
+
+## VRAM Requirements
+
+| Stage             | VRAM                   |
+| :---------------- | :--------------------- |
+| Stage-0 (AR)      | **~18 GiB + KV Cache** |
+| Stage-1 (DiT+VAE) | **~20 GiB**            |
+| Total             | **~38 GiB + KV Cache** |
+
+## File Description
+
+| File                        | Description                           |
+| --------------------------- | ------------------------------------- |
+| `run_server.sh`             | Server startup script                 |
+| `run_curl_text_to_image.sh` | Text-to-image curl example            |
+| `run_curl_image_edit.sh`    | Image-to-image (editing) curl example |
+| `openai_chat_client.py`     | Python client (t2i + i2i)             |
+
+## FAQ
+
+- If you encounter OOM errors, adjust `gpu_memory_utilization` in the stage config:
+
+```yaml
+# In glm_image.yaml, reduce from default 0.6:
+gpu_memory_utilization: 0.5
+```
+
+- The first request may be slow due to model warmup. Subsequent requests will be faster.
+
+## Example materials
+
+??? abstract "openai_chat_client.py"
+    ``````py
+    --8<-- "examples/online_serving/glm_image/openai_chat_client.py"
+    ``````
+??? abstract "run_curl_image_edit.sh"
+    ``````sh
+    --8<-- "examples/online_serving/glm_image/run_curl_image_edit.sh"
+    ``````
+??? abstract "run_curl_text_to_image.sh"
+    ``````sh
+    --8<-- "examples/online_serving/glm_image/run_curl_text_to_image.sh"
+    ``````
+??? abstract "run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/glm_image/run_server.sh"
+    ``````

docs/user_guide/examples/online_serving/image_to_video.md

@@ -0,0 +1,97 @@
+# Image-To-Video
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/image_to_video>.
+
+
+This example demonstrates how to deploy the Wan2.2 image-to-video model for online video generation using vLLM-Omni.
+
+## Start Server
+
+### Basic Start
+
+```bash
+vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --port 8091
+```
+
+### Start with Parameters
+
+Or use the startup script:
+
+```bash
+bash run_server.sh
+```
+
+The script allows overriding:
+- `MODEL` (default: `Wan-AI/Wan2.2-I2V-A14B-Diffusers`)
+- `PORT` (default: `8091`)
+- `BOUNDARY_RATIO` (default: `0.875`)
+- `FLOW_SHIFT` (default: `12.0`)
+- `CACHE_BACKEND` (default: `none`)
+- `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
+
+## API Calls
+
+### Method 1: Using curl
+
+```bash
+# Basic image-to-video generation
+bash run_curl_image_to_video.sh
+
+# Or execute directly (OpenAI-style multipart)
+curl -X POST http://localhost:8091/v1/videos \
+  -H "Accept: application/json" \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  -F "negative_prompt=low quality, blurry, static" \
+  -F "input_reference=@/path/to/qwen-bear.png" \
+  -F "width=832" \
+  -F "height=480" \
+  -F "num_frames=33" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=1.0" \
+  -F "guidance_scale_2=1.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=12.0" \
+  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_i2v_output.mp4
+```
+
+## Request Format
+
+### Required Fields
+
+```bash
+curl -X POST http://localhost:8091/v1/videos \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  -F "negative_prompt=low quality, blurry, static" \
+  -F "input_reference=@/path/to/qwen-bear.png"
+```
+
+### Generation with Parameters
+
+```bash
+curl -X POST http://localhost:8091/v1/videos \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  -F "negative_prompt=low quality, blurry, static" \
+  -F "input_reference=@/path/to/qwen-bear.png" \
+  -F "width=832" \
+  -F "height=480" \
+  -F "num_frames=33" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=1.0" \
+  -F "guidance_scale_2=1.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=12.0" \
+  -F "seed=42"
+```
+
+## Example materials
+
+??? abstract "run_curl_image_to_video.sh"
+    ``````sh
+    --8<-- "examples/online_serving/image_to_video/run_curl_image_to_video.sh"
+    ``````
+??? abstract "run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/image_to_video/run_server.sh"
+    ``````

docs/user_guide/examples/online_serving/qwen2_5_omni.md

@@ -30,7 +30,7 @@ cd examples/online_serving/qwen2_5_omni
 #### Send request via python
 
 ```bash
-python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
+python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities --port 8091 --host "localhost"
 ```
 
 The Python client supports the following command-line arguments:

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -36,7 +36,7 @@ cd examples/online_serving/qwen3_omni
 ####  Send request via python
 
 ```bash
-python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
+python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image --port 8091 --host "localhost"
 ```
 
 The Python client supports the following command-line arguments:

docs/user_guide/examples/online_serving/qwen3_tts.md

@@ -9,13 +9,13 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
 
 ## Supported Models
 
-| Model | Task Type | Description |
-|-------|-----------|-------------|
+| Model                                  | Task Type   | Description                                           |
+| -------------------------------------- | ----------- | ----------------------------------------------------- |
 | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
-| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
-| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
-| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
-| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description              |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-Base`        | Base        | Voice cloning from reference audio                    |
+| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant                                |
+| `Qwen/Qwen3-TTS-12Hz-0.6B-Base`        | Base        | Smaller/faster variant for voice cloning              |
 
 ## Run examples (Qwen3-TTS)
 
@@ -237,36 +237,58 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
 
 ### OpenAI Standard Parameters
 
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `input` | string | **required** | Text to synthesize |
-| `model` | string | server's model | Model to use (optional, should match server if specified) |
-| `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
-| `speed` | float | 1.0 | Playback speed (0.25-4.0) |
+| Parameter         | Type   | Default        | Description                                                 |
+| ----------------- | ------ | -------------- | ----------------------------------------------------------- |
+| `input`           | string | **required**   | Text to synthesize                                          |
+| `model`           | string | server's model | Model to use (optional, should match server if specified)   |
+| `voice`           | string | "vivian"       | Speaker name (e.g., vivian, ryan, aiden)                    |
+| `response_format` | string | "wav"          | Audio format: wav, mp3, flac, pcm, aac, opus                |
+| `speed`           | float  | 1.0            | Playback speed (0.25-4.0, not supported with `stream=true`) |
 
 ### vLLM-Omni Extension Parameters
 
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `task_type` | string | "CustomVoice" | Task: CustomVoice, VoiceDesign, or Base |
-| `language` | string | "Auto" | Language (see supported languages below) |
-| `instructions` | string | "" | Voice style/emotion instructions |
-| `max_new_tokens` | int | 2048 | Maximum tokens to generate |
+| Parameter        | Type   | Default       | Description                                                                  |
+| ---------------- | ------ | ------------- | ---------------------------------------------------------------------------- |
+| `task_type`      | string | "CustomVoice" | Task: CustomVoice, VoiceDesign, or Base                                      |
+| `language`       | string | "Auto"        | Language (see supported languages below)                                     |
+| `instructions`   | string | ""            | Voice style/emotion instructions                                             |
+| `max_new_tokens` | int    | 2048          | Maximum tokens to generate                                                   |
+| `stream`         | bool   | false         | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
 
 **Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
 
 ### Voice Clone Parameters (Base task)
 
-| Parameter | Type | Required | Description |
-|-----------|------|----------|-------------|
-| `ref_audio` | string | **Yes** | Reference audio (URL or base64 data URL) |
-| `ref_text` | string | No | Transcript of reference audio (for ICL mode) |
-| `x_vector_only_mode` | bool | No | Use speaker embedding only (no ICL) |
+| Parameter            | Type   | Required | Description                                  |
+| -------------------- | ------ | -------- | -------------------------------------------- |
+| `ref_audio`          | string | **Yes**  | Reference audio (URL or base64 data URL)     |
+| `ref_text`           | string | No       | Transcript of reference audio (for ICL mode) |
+| `x_vector_only_mode` | bool   | No       | Use speaker embedding only (no ICL)          |
+
+## Streaming
+
+Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded
+(one chunk per 25-frame Code2Wav window):
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, how are you?",
+        "voice": "vivian",
+        "language": "English",
+        "stream": true,
+        "response_format": "pcm"
+    }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
+```
+
+**Constraints:**
+- `stream=true` requires `response_format="pcm"` (raw 16-bit signed PCM, 24 kHz mono).
+- `speed` adjustment is not supported when streaming.
+- Requires the server stage config to have `async_chunk: true` (default in `qwen3_tts.yaml`).
 
 ## Limitations
 
-- **No streaming**: Audio is generated completely before being returned. Streaming will be supported after the pipeline is disaggregated (see RFC #938).
 - **Single request**: Batch processing is not yet optimized for online serving.
 
 ## Troubleshooting

docs/user_guide/examples/online_serving/text_to_video.md

@@ -0,0 +1,130 @@
+# Text-To-Video
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/text_to_video>.
+
+
+This example demonstrates how to deploy the Wan2.2 text-to-video model for online video generation using vLLM-Omni.
+
+## Start Server
+
+### Basic Start
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091
+```
+
+### Start with Parameters
+
+Or use the startup script:
+
+```bash
+bash run_server.sh
+```
+
+The script allows overriding:
+- `MODEL` (default: `Wan-AI/Wan2.2-T2V-A14B-Diffusers`)
+- `PORT` (default: `8091`)
+- `BOUNDARY_RATIO` (default: `0.875`)
+- `FLOW_SHIFT` (default: `5.0`)
+- `CACHE_BACKEND` (default: `none`)
+- `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
+
+## API Calls
+
+### Method 1: Using curl
+
+```bash
+# Basic text-to-video generation
+bash run_curl_text_to_video.sh
+
+# Or execute directly (OpenAI-style multipart)
+curl -s http://localhost:8091/v1/videos \
+  -H "Accept: application/json" \
+  -F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
+  -F "width=832" \
+  -F "height=480" \
+  -F "num_frames=33" \
+  -F "negative_prompt=色调艳丽 ，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=4.0" \
+  -F "guidance_scale_2=4.0" \
+  -F "boundary_ratio=0.875" \
+  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
+```
+
+## Request Format
+
+### Simple Text-to-Video Generation
+
+```bash
+curl -X POST http://localhost:8091/v1/videos \
+  -F "prompt=A cinematic view of a futuristic city at sunset"
+```
+
+### Generation with Parameters
+
+```bash
+curl -X POST http://localhost:8091/v1/videos \
+  -F "prompt=A cinematic view of a futuristic city at sunset" \
+  -F "width=832" \
+  -F "height=480" \
+  -F "num_frames=33" \
+  -F "negative_prompt=low quality, blurry, static" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=4.0" \
+  -F "guidance_scale_2=4.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=5.0" \
+  -F "seed=42"
+```
+
+## Generation Parameters
+
+| Parameter             | Type   | Default | Description                                      |
+| --------------------- | ------ | ------- | ------------------------------------------------ |
+| `prompt`              | str    | -       | Text description of the desired video            |
+| `negative_prompt`     | str    | None    | Negative prompt                                  |
+| `n`                   | int    | 1       | Number of videos to generate                     |
+| `width`               | int    | None    | Video width in pixels                            |
+| `height`              | int    | None    | Video height in pixels                           |
+| `num_frames`          | int    | None    | Number of frames to generate                     |
+| `fps`                 | int    | None    | Frames per second for output video               |
+| `num_inference_steps` | int    | None    | Number of denoising steps                        |
+| `guidance_scale`      | float  | None    | CFG guidance scale (low-noise stage)             |
+| `guidance_scale_2`    | float  | None    | CFG guidance scale (high-noise stage, Wan2.2)     |
+| `boundary_ratio`      | float  | None    | Boundary split ratio for low/high DiT (Wan2.2)   |
+| `flow_shift`          | float  | None    | Scheduler flow shift (Wan2.2)                    |
+| `seed`                | int    | None    | Random seed (reproducible)                       |
+| `lora`                | object | None    | LoRA configuration                               |
+| `extra_body`          | object | None    | Model-specific extra parameters                  |
+
+## Response Format
+
+```json
+{
+  "created": 1234567890,
+  "data": [
+    { "b64_json": "<base64-mp4>" }
+  ]
+}
+```
+
+## Extract Video
+
+```bash
+# Extract base64 from response and decode to video
+cat response.json | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
+```
+
+## Example materials
+
+??? abstract "run_curl_text_to_video.sh"
+    ``````sh
+    --8<-- "examples/online_serving/text_to_video/run_curl_text_to_video.sh"
+    ``````
+??? abstract "run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/text_to_video/run_server.sh"
+    ``````

examples/offline_inference/glm_image/README.md

@@ -0,0 +1,138 @@
+# GLM-Image Multistage End-to-End Inference
+
+This example demonstrates how to run GLM-Image with the vLLM-Omni multistage architecture.
+
+## Architecture
+
+GLM-Image uses a 2-stage pipeline:
+
+```
+┌─────────────────────────────────────────────────────────────┐
+│                     GLM-Image Pipeline                       │
+├─────────────────────────────────────────────────────────────┤
+│                                                             │
+│  Stage 0 (AR Model)              Stage 1 (Diffusion)        │
+│  ┌─────────────────┐            ┌─────────────────────┐     │
+│  │ vLLM-optimized  │            │  GlmImagePipeline   │     │
+│  │ GlmImageFor     │  prior     │  ┌───────────────┐  │     │
+│  │ Conditional     │──tokens───►│  │ DiT Denoiser  │  │     │
+│  │ Generation      │            │  └───────────────┘  │     │
+│  │ (9B AR model)   │            │         │          │     │
+│  └─────────────────┘            │         ▼          │     │
+│         ▲                       │  ┌───────────────┐  │     │
+│         │                       │  │  VAE Decode   │──┼──► Image
+│    Text/Image                   │  └───────────────┘  │     │
+│      Input                      └─────────────────────┘     │
+│                                                             │
+└─────────────────────────────────────────────────────────────┘
+```
+
+## Features
+
+- **vLLM-optimized AR**: Uses PagedAttention and tensor parallelism for faster prior token generation
+- **Flexible deployment**: AR and Diffusion stages can run on different GPUs
+- **Text-to-Image**: Generate images from text descriptions
+- **Image-to-Image**: Edit existing images with text prompts
+
+## Usage
+
+### Text-to-Image
+
+```bash
+python end2end.py \
+    --model-path /path/to/glm-image \
+    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --prompt "A beautiful sunset over the ocean with sailing boats" \
+    --height 1024 \
+    --width 1024 \
+    --output output_t2i.png
+```
+
+### Image-to-Image (Image Editing)
+
+```bash
+python end2end.py \
+    --model-path /path/to/glm-image \
+    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --prompt "Transform this scene into a winter wonderland" \
+    --image input.png \
+    --output output_i2i.png
+```
+
+### With Custom Parameters
+
+```bash
+python end2end.py \
+    --model-path /path/to/glm-image \
+    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --prompt "A photorealistic cat sitting on a window sill" \
+    --height 1024 \
+    --width 1024 \
+    --num-inference-steps 50 \
+    --guidance-scale 1.5 \
+    --seed 42 \
+    --output output.png
+```
+
+## Shell Scripts
+
+### Run Text-to-Image
+
+```bash
+./run_t2i.sh
+```
+
+### Run Image-to-Image
+
+```bash
+./run_i2i.sh --image /path/to/input.png
+```
+
+## Stage Configuration
+
+The stage config (`glm_image.yaml`) defines:
+
+- **Stage 0 (AR)**: Uses `GPUARWorker` with vLLM engine
+
+  - Model: `GlmImageForConditionalGeneration`
+  - Output: `token_ids` (prior tokens)
+
+- **Stage 1 (Diffusion)**: Uses diffusion engine
+  - Model: `GlmImagePipeline`
+  - Output: Generated image
+
+See `vllm_omni/model_executor/stage_configs/glm_image.yaml` for full configuration.
+
+## Comparison with Single-Stage
+
+| Aspect      | Single-Stage (transformers) | Multistage (vLLM)   |
+| ----------- | --------------------------- | ------------------- |
+| AR Model    | transformers native         | vLLM PagedAttention |
+| Memory      | Higher (no KV cache opt)    | Lower (optimized)   |
+| Throughput  | Lower                       | Higher              |
+| Flexibility | Single GPU                  | Multi-GPU support   |
+
+## Troubleshooting
+
+### OOM Error
+
+Try reducing memory usage:
+
+```bash
+# In glm_image.yaml, adjust:
+gpu_memory_utilization: 0.5  # Reduce from 0.6
+```
+
+### Slow Initialization
+
+The first run loads model weights. Subsequent runs are faster:
+
+```bash
+--stage-init-timeout 900  # Increase timeout for slow storage
+```
+
+## Requirements
+
+- vLLM-Omni with GLM-Image support
+- CUDA-capable GPU (recommended: H100/A100 with 80GB)
+- GLM-Image model weights

examples/offline_inference/glm_image/end2end.py

@@ -0,0 +1,476 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+End-to-end offline inference example for GLM-Image with multistage architecture.
+
+This script tests the multistage pipeline where:
+- Stage 0 (AR): vLLM-optimized GlmImageForConditionalGeneration generates prior_token_ids
+- Stage 1 (Diffusion): GlmImagePipeline performs DiT denoising + VAE decode
+
+Usage (text-to-image):
+    python end2end.py \
+        --model-path /path/to/glm-image \
+        --config-path /path/to/glm_image.yaml \
+        --prompt "A beautiful sunset over the ocean" \
+        --output output_t2i.png
+
+Usage (image-to-image / image edit):
+    python end2end.py \
+        --model-path /path/to/glm-image \
+        --config-path /path/to/glm_image.yaml \
+        --prompt "Make it look like winter" \
+        --image input.png \
+        --output output_i2i.png
+
+Usage (with custom parameters):
+    python end2end.py \
+        --model-path /path/to/glm-image \
+        --config-path /path/to/glm_image.yaml \
+        --prompt "A cat sitting on a window sill" \
+        --height 1024 \
+        --width 1024 \
+        --num-inference-steps 50 \
+        --guidance-scale 1.5 \
+        --seed 42
+
+For more options, run:
+    python end2end.py --help
+"""
+
+import argparse
+import os
+import time
+from pathlib import Path
+
+from PIL import Image
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+# Default stage config path (relative to vllm_omni package)
+DEFAULT_CONFIG_PATH = "vllm_omni/model_executor/stage_configs/glm_image.yaml"
+
+SEED = 42
+
+# GLM-Image special tokens
+GLM_IMAGE_EOS_TOKEN_ID = 16385  # eos_token_id from generation_config.json
+GLM_IMAGE_VISION_VOCAB_SIZE = 16512  # top_k should be vision_vocab_size
+
+
+def compute_max_tokens(height: int, width: int, factor: int = 32) -> int:
+    """
+    Compute max_new_tokens for GLM-Image AR generation.
+
+    GLM-Image generates tokens in this order for text-to-image:
+    1. Small preview image (half resolution in each dimension)
+    2. Large target image (full resolution)
+    3. EOS token
+
+    Args:
+        height: Target image height in pixels
+        width: Target image width in pixels
+        factor: Downsampling factor (32 for GLM-Image AR output)
+
+    Returns:
+        Total number of tokens to generate (small + large + EOS)
+    """
+    # Large image tokens (target resolution)
+    token_h = height // factor
+    token_w = width // factor
+    large_tokens = token_h * token_w
+
+    # Small preview tokens (half resolution in each dimension)
+    small_h = token_h // 2
+    small_w = token_w // 2
+    small_tokens = small_h * small_w
+
+    # Total: small + large + EOS
+    return small_tokens + large_tokens + 1
+
+
+def load_image(image_path: str) -> Image.Image:
+    """Load an image from file path."""
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(f"Image file not found: {image_path}")
+    return Image.open(image_path).convert("RGB")
+
+
+def save_image(image: Image.Image, output_path: str) -> None:
+    """Save an image to file path."""
+    output_dir = os.path.dirname(output_path)
+    if output_dir:
+        os.makedirs(output_dir, exist_ok=True)
+    image.save(output_path)
+    print(f"Image saved to: {output_path}")
+
+
+def build_prompt_for_t2i(
+    prompt: str,
+    height: int = 1024,
+    width: int = 1024,
+) -> dict:
+    """
+    Build prompt dict for text-to-image generation.
+
+    Args:
+        prompt: Text description for image generation
+        height: Target image height
+        width: Target image width
+
+    Returns:
+        Dict containing prompt and generation parameters
+    """
+    return {
+        "prompt": prompt,
+        "height": height,
+        "width": width,
+        # Pass target dimensions to AR processor for proper grid token generation
+        "mm_processor_kwargs": {
+            "target_h": height,
+            "target_w": width,
+        },
+    }
+
+
+def build_prompt_for_i2i(
+    prompt: str,
+    image: Image.Image,
+    height: int | None = None,
+    width: int | None = None,
+) -> dict:
+    """
+    Build prompt dict for image-to-image generation.
+
+    Args:
+        prompt: Text description for image editing
+        image: Source image for editing
+        height: Target image height (default: use source image size)
+        width: Target image width (default: use source image size)
+
+    Returns:
+        Dict containing prompt, image, and generation parameters
+    """
+    # Use source image dimensions if not specified
+    if height is None:
+        height = image.height
+    if width is None:
+        width = image.width
+
+    return {
+        "prompt": prompt,
+        "multi_modal_data": {
+            "image": image,
+        },
+        "height": height,
+        "width": width,
+        # Pass target dimensions to AR processor for proper grid token generation
+        "mm_processor_kwargs": {
+            "target_h": height,
+            "target_w": width,
+        },
+    }
+
+
+def main(args: argparse.Namespace) -> None:
+    """Main entry point for GLM-Image end-to-end inference."""
+    print("=" * 60)
+    print("GLM-Image Multistage End-to-End Inference")
+    print("=" * 60)
+
+    # Validate arguments
+    if not args.model_path:
+        raise ValueError("--model-path is required")
+
+    if not args.prompt:
+        raise ValueError("--prompt is required")
+
+    # Determine config path
+    config_path = args.config_path
+    if config_path is None:
+        # Try to find default config
+        if os.path.exists(DEFAULT_CONFIG_PATH):
+            config_path = DEFAULT_CONFIG_PATH
+        else:
+            # Try relative to script location
+            script_dir = Path(__file__).parent.parent.parent.parent
+            config_path = script_dir / "vllm_omni/model_executor/stage_configs/glm_image.yaml"
+            if not config_path.exists():
+                raise FileNotFoundError(
+                    f"Stage config not found. Please specify --config-path. Tried: {DEFAULT_CONFIG_PATH}"
+                )
+            config_path = str(config_path)
+
+    print(f"Model path: {args.model_path}")
+    print(f"Config path: {config_path}")
+    print(f"Prompt: {args.prompt}")
+
+    # Load source image for image-to-image mode
+    source_image = None
+    if args.image:
+        print(f"Source image: {args.image}")
+        source_image = load_image(args.image)
+        print(f"  Image size: {source_image.width}x{source_image.height}")
+
+    # Build prompt based on mode
+    if source_image is not None:
+        # Image-to-image mode
+        prompt_dict = build_prompt_for_i2i(
+            prompt=args.prompt,
+            image=source_image,
+            height=args.height,
+            width=args.width,
+        )
+        mode = "image-to-image"
+    else:
+        # Text-to-image mode
+        prompt_dict = build_prompt_for_t2i(
+            prompt=args.prompt,
+            height=args.height or 1024,
+            width=args.width or 1024,
+        )
+        mode = "text-to-image"
+
+    print(f"Mode: {mode}")
+    print(f"Target size: {prompt_dict.get('height', 1024)}x{prompt_dict.get('width', 1024)}")
+
+    # Add generation parameters to prompt
+    prompt_dict["seed"] = args.seed
+    prompt_dict["num_inference_steps"] = args.num_inference_steps
+    prompt_dict["guidance_scale"] = args.guidance_scale
+
+    if args.negative_prompt:
+        prompt_dict["negative_prompt"] = args.negative_prompt
+
+    # Initialize Omni with multistage config
+    print("\nInitializing Omni with multistage pipeline...")
+    start_time = time.time()
+
+    omni = Omni(
+        model=args.model_path,
+        stage_configs_path=config_path,
+        log_stats=args.enable_stats,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+
+    init_time = time.time() - start_time
+    print(f"Initialization completed in {init_time:.2f}s")
+
+    # Prepare prompts (support batch generation)
+    prompts = [prompt_dict for _ in range(args.num_prompts)]
+
+    # No explicit sampling_params for diffusion - parameters are in prompt_dict
+    # For multistage, the AR stage may need sampling params
+    from vllm import SamplingParams
+
+    # Compute max_tokens dynamically based on target image size
+    target_height = prompt_dict.get("height", 1024)
+    target_width = prompt_dict.get("width", 1024)
+    calculated_max_tokens = compute_max_tokens(target_height, target_width)
+
+    # Use calculated value unless user explicitly specified a different value
+    # Default args.max_tokens is 16384 (very large), so prefer calculated value
+    effective_max_tokens = calculated_max_tokens if args.max_tokens == 16384 else args.max_tokens
+
+    if args.verbose:
+        print(f"AR max_tokens: {effective_max_tokens} (calculated: {calculated_max_tokens}, arg: {args.max_tokens})")
+
+    # IMPORTANT: GLM-Image AR model requires these exact sampling parameters
+    # from generation_config.json for proper image token generation.
+    # - temperature=0.9, top_p=0.75, top_k=16512 (vision_vocab_size)
+    # - stop_token_ids=[16385] (eos_token_id) is CRITICAL to stop generation
+    ar_sampling_params = SamplingParams(
+        temperature=0.9,  # From generation_config.json
+        top_p=0.75,  # From generation_config.json
+        top_k=GLM_IMAGE_VISION_VOCAB_SIZE,  # 16512, vision vocabulary size
+        max_tokens=effective_max_tokens,
+        stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID],  # 16385, CRITICAL for stopping
+        seed=args.seed,
+        detokenize=False,
+    )
+
+    # For diffusion stage, sampling_params contains diffusion-specific parameters
+    # These are passed as kwargs to the diffusion engine
+    diffusion_sampling_params = OmniDiffusionSamplingParams(
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+        height=prompt_dict.get("height", 1024),
+        width=prompt_dict.get("width", 1024),
+        seed=args.seed,
+    )
+
+    # For multistage, we need sampling_params for each stage
+    # Stage 0 (AR): SamplingParams for vLLM
+    # Stage 1 (Diffusion): dict with diffusion kwargs
+    sampling_params_list = [ar_sampling_params, diffusion_sampling_params]
+
+    # Run generation
+    print(f"\nGenerating {args.num_prompts} image(s)...")
+    gen_start_time = time.time()
+
+    output_dir = os.path.dirname(args.output) if args.output else "outputs"
+    if output_dir:
+        os.makedirs(output_dir, exist_ok=True)
+
+    output_count = 0
+    for stage_outputs in omni.generate(prompts, sampling_params_list, py_generator=True):
+        if stage_outputs.final_output_type == "image":
+            for output in stage_outputs.request_output:
+                request_id = output.request_id
+
+                # Get generated images
+                images = output.images if hasattr(output, "images") else []
+                if not images and hasattr(output, "multimodal_output"):
+                    images = output.multimodal_output.get("images", [])
+
+                # Save each generated image
+                for idx, img in enumerate(images):
+                    if args.num_prompts == 1 and len(images) == 1:
+                        output_path = args.output
+                    else:
+                        base, ext = os.path.splitext(args.output)
+                        output_path = f"{base}_{request_id}_{idx}{ext}"
+
+                    if isinstance(img, Image.Image):
+                        save_image(img, output_path)
+                    else:
+                        print(f"Warning: Unexpected image type for request {request_id}: {type(img)}")
+
+                    output_count += 1
+
+        elif stage_outputs.final_output_type == "text":
+            # AR stage output (intermediate, for debugging)
+            if args.verbose:
+                for output in stage_outputs.request_output:
+                    print(f"AR output for request {output.request_id}:")
+                    print(f"  Token count: {len(output.outputs[0].token_ids)}")
+
+    gen_time = time.time() - gen_start_time
+    print(f"\nGeneration completed in {gen_time:.2f}s")
+    print(f"Generated {output_count} image(s)")
+
+    # Cleanup
+    omni.close()
+    print("\nDone!")
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="GLM-Image Multistage End-to-End Inference",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+
+    # Required arguments
+    parser.add_argument(
+        "--model-path",
+        type=str,
+        required=True,
+        help="Path to GLM-Image model directory or HuggingFace model ID",
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        required=True,
+        help="Text prompt for image generation",
+    )
+
+    # Optional arguments
+    parser.add_argument(
+        "--config-path",
+        type=str,
+        default=None,
+        help="Path to stage config YAML file (default: auto-detect)",
+    )
+    parser.add_argument(
+        "--image",
+        type=str,
+        default=None,
+        help="Path to source image for image-to-image mode",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="output_glm_image.png",
+        help="Output image path (default: output_glm_image.png)",
+    )
+    parser.add_argument(
+        "--negative-prompt",
+        type=str,
+        default=None,
+        help="Negative prompt for classifier-free guidance",
+    )
+
+    # Generation parameters
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=None,
+        help="Output image height (default: 1024 for t2i, source size for i2i)",
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=None,
+        help="Output image width (default: 1024 for t2i, source size for i2i)",
+    )
+    parser.add_argument(
+        "--num-inference-steps",
+        type=int,
+        default=50,
+        help="Number of diffusion denoising steps (default: 50)",
+    )
+    parser.add_argument(
+        "--guidance-scale",
+        type=float,
+        default=1.5,
+        help="Classifier-free guidance scale (default: 1.5)",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=SEED,
+        help=f"Random seed for reproducibility (default: {SEED})",
+    )
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=16384,
+        help="Maximum tokens for AR generation (default: 16384)",
+    )
+
+    # Batch processing
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=1,
+        help="Number of images to generate (default: 1)",
+    )
+
+    # Runtime options
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        default=False,
+        help="Enable statistics logging",
+    )
+    parser.add_argument(
+        "--stage-init-timeout",
+        type=int,
+        default=600,
+        help="Timeout for stage initialization in seconds (default: 600)",
+    )
+    parser.add_argument(
+        "--verbose",
+        "-v",
+        action="store_true",
+        default=False,
+        help="Enable verbose output",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

examples/offline_inference/glm_image/run_i2i.sh

@@ -0,0 +1,93 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# Run GLM-Image image-to-image (editing) with multistage pipeline
+
+set -e
+
+# Default values
+MODEL_PATH="${MODEL_PATH:-/path/to/glm-image}"
+CONFIG_PATH="${CONFIG_PATH:-vllm_omni/model_executor/stage_configs/glm_image.yaml}"
+PROMPT="${PROMPT:-Transform this image into an oil painting style}"
+INPUT_IMAGE=""
+OUTPUT="${OUTPUT:-output_i2i.png}"
+NUM_STEPS="${NUM_STEPS:-50}"
+GUIDANCE="${GUIDANCE:-1.5}"
+SEED="${SEED:-42}"
+
+# Parse command line arguments
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --model-path)
+            MODEL_PATH="$2"
+            shift 2
+            ;;
+        --config-path)
+            CONFIG_PATH="$2"
+            shift 2
+            ;;
+        --prompt)
+            PROMPT="$2"
+            shift 2
+            ;;
+        --image)
+            INPUT_IMAGE="$2"
+            shift 2
+            ;;
+        --output)
+            OUTPUT="$2"
+            shift 2
+            ;;
+        --num-steps)
+            NUM_STEPS="$2"
+            shift 2
+            ;;
+        --guidance)
+            GUIDANCE="$2"
+            shift 2
+            ;;
+        --seed)
+            SEED="$2"
+            shift 2
+            ;;
+        *)
+            echo "Unknown option: $1"
+            exit 1
+            ;;
+    esac
+done
+
+# Check if input image is provided
+if [ -z "${INPUT_IMAGE}" ]; then
+    echo "Error: --image is required for image-to-image mode"
+    echo "Usage: ./run_i2i.sh --image /path/to/input.png [--prompt \"edit instruction\"]"
+    exit 1
+fi
+
+if [ ! -f "${INPUT_IMAGE}" ]; then
+    echo "Error: Input image not found: ${INPUT_IMAGE}"
+    exit 1
+fi
+
+echo "=============================================="
+echo "GLM-Image Image-to-Image Generation"
+echo "=============================================="
+echo "Model: ${MODEL_PATH}"
+echo "Config: ${CONFIG_PATH}"
+echo "Input: ${INPUT_IMAGE}"
+echo "Prompt: ${PROMPT}"
+echo "Output: ${OUTPUT}"
+echo "Steps: ${NUM_STEPS}"
+echo "Guidance: ${GUIDANCE}"
+echo "Seed: ${SEED}"
+echo "=============================================="
+
+python end2end.py \
+    --model-path "${MODEL_PATH}" \
+    --config-path "${CONFIG_PATH}" \
+    --prompt "${PROMPT}" \
+    --image "${INPUT_IMAGE}" \
+    --output "${OUTPUT}" \
+    --num-inference-steps "${NUM_STEPS}" \
+    --guidance-scale "${GUIDANCE}" \
+    --seed "${SEED}" \
+    --verbose

examples/offline_inference/glm_image/run_t2i.sh

@@ -0,0 +1,87 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# Run GLM-Image text-to-image generation with multistage pipeline
+
+set -e
+
+# Default values
+MODEL_PATH="${MODEL_PATH:-/path/to/glm-image}"
+CONFIG_PATH="${CONFIG_PATH:-vllm_omni/model_executor/stage_configs/glm_image.yaml}"
+PROMPT="${PROMPT:-A beautiful mountain landscape with snow-capped peaks and a clear blue lake}"
+OUTPUT="${OUTPUT:-output_t2i.png}"
+HEIGHT="${HEIGHT:-1024}"
+WIDTH="${WIDTH:-1024}"
+NUM_STEPS="${NUM_STEPS:-50}"
+GUIDANCE="${GUIDANCE:-1.5}"
+SEED="${SEED:-42}"
+
+# Parse command line arguments
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --model-path)
+            MODEL_PATH="$2"
+            shift 2
+            ;;
+        --config-path)
+            CONFIG_PATH="$2"
+            shift 2
+            ;;
+        --prompt)
+            PROMPT="$2"
+            shift 2
+            ;;
+        --output)
+            OUTPUT="$2"
+            shift 2
+            ;;
+        --height)
+            HEIGHT="$2"
+            shift 2
+            ;;
+        --width)
+            WIDTH="$2"
+            shift 2
+            ;;
+        --num-steps)
+            NUM_STEPS="$2"
+            shift 2
+            ;;
+        --guidance)
+            GUIDANCE="$2"
+            shift 2
+            ;;
+        --seed)
+            SEED="$2"
+            shift 2
+            ;;
+        *)
+            echo "Unknown option: $1"
+            exit 1
+            ;;
+    esac
+done
+
+echo "=============================================="
+echo "GLM-Image Text-to-Image Generation"
+echo "=============================================="
+echo "Model: ${MODEL_PATH}"
+echo "Config: ${CONFIG_PATH}"
+echo "Prompt: ${PROMPT}"
+echo "Output: ${OUTPUT}"
+echo "Size: ${WIDTH}x${HEIGHT}"
+echo "Steps: ${NUM_STEPS}"
+echo "Guidance: ${GUIDANCE}"
+echo "Seed: ${SEED}"
+echo "=============================================="
+
+python end2end.py \
+    --model-path "${MODEL_PATH}" \
+    --config-path "${CONFIG_PATH}" \
+    --prompt "${PROMPT}" \
+    --output "${OUTPUT}" \
+    --height "${HEIGHT}" \
+    --width "${WIDTH}" \
+    --num-inference-steps "${NUM_STEPS}" \
+    --guidance-scale "${GUIDANCE}" \
+    --seed "${SEED}" \
+    --verbose

examples/offline_inference/image_to_video/README.md

@@ -66,5 +66,10 @@ Key arguments:
 - `--vae-use-tiling`: Enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+- `--use-hsdp`: Enable Hybrid Sharded Data Parallel to shard model weights across GPUs.
+- `--hsdp-shard-size`: Number of GPUs to shard model weights across within each replica group. -1 (default) auto-calculates as world_size / replicate_size.
+- `--hsdp-replicate-size`: Number of replica groups for HSDP. Each replica holds a full sharded copy. Default 1 means pure sharding (no replication).
+
+
 
 > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.

examples/offline_inference/image_to_video/image_to_video.py

@@ -80,6 +80,24 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument(
+        "--ulysses-degree",
+        type=int,
+        default=1,
+        help="Number of GPUs used for ulysses sequence parallelism.",
+    )
+    parser.add_argument(
+        "--ring-degree",
+        type=int,
+        default=1,
+        help="Number of GPUs used for ring sequence parallelism.",
+    )
+    parser.add_argument(
+        "--tensor-parallel-size",
+        type=int,
+        default=1,
+        help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
+    )
     parser.add_argument(
         "--cfg-parallel-size",
         type=int,
@@ -92,6 +110,29 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Disable torch.compile and force eager execution.",
     )
+    parser.add_argument(
+        "--use-hsdp",
+        action="store_true",
+        help=("Enable Hybrid Sharded Data Parallel to shard model weights across GPUs. "),
+    )
+    parser.add_argument(
+        "--hsdp-shard-size",
+        type=int,
+        default=-1,
+        help=(
+            "Number of GPUs to shard model weights across within each replica group. "
+            "-1 (default) auto-calculates as world_size / replicate_size. "
+        ),
+    )
+    parser.add_argument(
+        "--hsdp-replicate-size",
+        type=int,
+        default=1,
+        help=(
+            "Number of replica groups for HSDP. Each replica holds a full sharded copy. "
+            "Default 1 means pure sharding (no replication). "
+        ),
+    )
     return parser.parse_args()
 
 
@@ -128,7 +169,13 @@ def main():
     # Check if profiling is requested via environment variable
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
     parallel_config = DiffusionParallelConfig(
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
         cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
+        use_hsdp=args.use_hsdp,
+        hsdp_shard_size=args.hsdp_shard_size,
+        hsdp_replicate_size=args.hsdp_replicate_size,
     )
     omni = Omni(
         model=args.model,

examples/offline_inference/mimo_audio/README.md

@@ -0,0 +1,232 @@
+# MiMo-Audio Offline Inference
+
+This directory contains an offline demo for running MiMo-Audio models with vLLM Omni. It builds task-specific inputs and generates WAV files or text outputs locally.
+
+## Model Overview
+
+MiMo-Audio provides multiple task variants for audio understanding and generation:
+
+- **tts_sft**: Basic text-to-speech generation from text input.
+- **tts_sft_with_instruct**: TTS generation with explicit voice style instructions.
+- **tts_sft_with_audio**: TTS generation with audio reference for voice cloning.
+- **tts_sft_with_natural_instruction**: TTS generation from natural language descriptions embedded in text.
+- **audio_trancribing_sft**: Transcribe audio to text (speech-to-text).
+- **audio_understanding_sft**: Understand and analyze audio content with text queries.
+- **audio_understanding_sft_with_thinking**: Audio understanding with reasoning chain.
+- **spoken_dialogue_sft_multiturn**: Multi-turn spoken dialogue with audio input/output.
+- **speech2text_dialogue_sft_multiturn**: Multi-turn dialogue converting speech to text.
+- **text_dialogue_sft_multiturn**: Multi-turn text-only dialogue.
+
+## Setup
+
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
+
+### Environment Variables
+
+The `MIMO_AUDIO_TOKENIZER_PATH` environment variable is mandatory due to the specialized architecture:
+
+```bash
+export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
+```
+
+## Quick Start
+
+Run a single sample for basic TTS:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type tts_sft
+```
+
+Run batch samples for basic TTS:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type tts_sft \
+  --num-prompts {batch_size}
+```
+
+When enabling multi-batch processing, if the total number of tokens passed to the next stage exceeds the `max_model_len` value in the `mimo_audio.yaml` configuration file, you must also synchronously update the `max_position_embeddings` value in `MiMo-Audio-7B-Instruct/config.json` to match the modified value.
+
+Generated audio files are saved to `output_audio/` by default. `--num-prompts` also can be used to all tasks below.
+
+## Task Usage
+
+### tts_sft (Basic Text-to-Speech)
+
+Generate speech from text input:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type tts_sft \
+  --text "The weather is so nice today."
+```
+
+### tts_sft_with_instruct (TTS with Voice Instructions)
+
+Generate speech with explicit voice style instructions:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type tts_sft_with_instruct \
+  --text "The weather is so nice today." \
+  --instruct "Speak happily in a child's voice"
+```
+
+### tts_sft_with_audio (TTS with Audio Reference)
+
+Generate speech using an audio reference for voice cloning:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type tts_sft_with_audio \
+  --text "The weather is so nice today." \
+  --audio-path "./spoken_dialogue_assistant_turn_1.wav"
+```
+
+### tts_sft_with_natural_instruction (Natural Language TTS)
+
+Generate speech from text containing natural voice descriptions:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type tts_sft_with_natural_instruction \
+  --text "In a panting young male voice, he said: I can't run anymore, wait for me!"
+```
+
+### audio_trancribing_sft (Speech-to-Text)
+
+Transcribe audio to text:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type audio_trancribing_sft \
+  --audio-path "./spoken_dialogue_assistant_turn_1.wav"
+```
+
+### audio_understanding_sft (Audio Understanding)
+
+Understand and analyze audio content with text queries:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type audio_understanding_sft \
+  --text "Summarize the audio." \
+  --audio-path "./spoken_dialogue_assistant_turn_1.wav"
+```
+
+### audio_understanding_sft_with_thinking (Audio Understanding with Reasoning)
+
+Audio understanding with reasoning chain:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type audio_understanding_sft_with_thinking \
+  --text "Summarize the audio." \
+  --audio-path "./spoken_dialogue_assistant_turn_1.wav"
+```
+
+### spoken_dialogue_sft_multiturn (Multi-turn Spoken Dialogue)
+
+Multi-turn dialogue with audio input and output:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type spoken_dialogue_sft_multiturn \
+  --audio-path "./prompt_speech_zh_m.wav"
+```
+
+Note: This task uses hardcoded audio files in the script. The audio files used in examples are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples
+
+### speech2text_dialogue_sft_multiturn (Speech-to-Text Dialogue)
+
+Multi-turn dialogue converting speech to text:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type speech2text_dialogue_sft_multiturn
+```
+
+Note: This task uses hardcoded audio files and message lists in the script.
+
+### text_dialogue_sft_multiturn (Text Dialogue)
+
+Multi-turn text-only dialogue:
+
+```bash
+python3 -u end2end.py \
+  --stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+  --model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
+  --query-type text_dialogue_sft_multiturn
+```
+
+Note: This task uses hardcoded message lists in the script.
+
+## Troubleshooting
+
+### Audio dependencies (soundfile, librosa)
+
+This example depends on **soundfile** (read/write WAV) and **librosa** (load audio including MP3). Install the project requirements first:
+
+```bash
+pip install -r requirements/common.txt
+# or at least: pip install soundfile>=0.13.1 librosa>=0.11.0
+```
+
+- **`soundfile` / libsndfile not found**  
+  `soundfile` uses the C library **libsndfile**. On Linux, install the system package before pip:
+  - Debian/Ubuntu: `sudo apt-get install libsndfile1`
+  - For development builds: `sudo apt-get install libsndfile1-dev`
+  - Then: `pip install soundfile`
+
+- **`librosa` fails to load MP3 or reports "No backend available"**  
+  Loading MP3 (e.g. in `spoken_dialogue_sft_multiturn` with `.mp3` files) uses **ffmpeg** as the backend. Install ffmpeg:
+  - Debian/Ubuntu: `sudo apt-get install ffmpeg`
+  - macOS: `brew install ffmpeg`
+
+- **`ImportError: No module named 'soundfile'` or `ModuleNotFoundError: ... librosa`**  
+  Ensure you are in the same Python environment where vLLM Omni and the example dependencies are installed, and that `requirements/common.txt` (or the packages above) are installed.
+
+### Tokenizer path
+
+- **`MIMO_AUDIO_TOKENIZER_PATH` not set or model fails to find tokenizer**  
+  Export the tokenizer path before running:
+  ```bash
+  export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
+  ```
+  See [Environment Variables](#environment-variables) in Setup.
+
+### Other
+
+- If the model or stage config fails to load, check [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for memory and GPU settings.
+- For errors when reading/writing WAV (e.g. unsupported format), ensure input files are standard WAV/MP3 and that `soundfile` is linked to a working libsndfile (see above).
+
+## Notes
+
+- The script uses default model paths and audio files embedded in `end2end.py`. Update them if your local cache path differs.
+- Use `--output-dir` to change the output folder (default: `./output_audio`).
+- Use `--num-prompts` to generate multiple prompts in one run (default: 1).
+- Audio files used in multi-turn dialogue examples are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples
+- The script supports various configuration options for initialization timeouts, batch timeouts, and shared memory thresholds. See `--help` for details.

examples/offline_inference/mimo_audio/end2end.py

@@ -0,0 +1,443 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+This example shows how to use vLLM for running offline inference
+with the correct prompt format on MiMo-Audio-Omni.
+"""
+
+import copy
+import json
+import os
+from typing import NamedTuple
+
+import soundfile as sf
+from message_convert import (
+    get_audio_data,
+    get_audio_understanding_sft_prompt,
+    get_s2t_dialogue_sft_multiturn_prompt,
+    get_spoken_dialogue_sft_multiturn_prompt,
+    get_text_dialogue_sft_multiturn_prompt,
+    get_tts_sft_prompt,
+    to_prompt,
+)
+from vllm import SamplingParams
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniTokensPrompt
+
+SEED = 42
+
+MAX_CODE2WAV_TOKENS = 18192  # Maximum tokens supported by code2wav model
+
+
+class QueryResult(NamedTuple):
+    inputs: dict
+    limit_mm_per_prompt: dict[str, int]
+
+
+def get_codes_query_from_json(codes_path: str) -> QueryResult:
+    with open(codes_path, encoding="utf-8") as f:
+        data = json.load(f)
+
+    if isinstance(data, list):
+        code_final = data
+    elif isinstance(data, dict) and "code_final" in data:
+        code_final = data["code_final"]
+    else:
+        raise ValueError(
+            f"Unsupported codes json format in {codes_path}.\n"
+            "Expect a JSON list[int] or {{'code_final': list[int]}}."
+        )
+
+    if not isinstance(code_final, list) or not all(isinstance(x, int) for x in code_final):
+        raise ValueError("code_final must be a list[int].")
+
+    if len(code_final) > MAX_CODE2WAV_TOKENS:
+        print(f"[Warn] code_final len={len(code_final)} > {MAX_CODE2WAV_TOKENS}, truncating.")
+        code_final = code_final[:MAX_CODE2WAV_TOKENS]
+
+    return QueryResult(
+        inputs=OmniTokensPrompt(
+            prompt_token_ids=code_final,
+            multi_modal_data=None,
+            mm_processor_kwargs=None,
+        ),
+        limit_mm_per_prompt={},
+    )
+
+
+def get_tts_sft(
+    text="The weather is so nice today.",
+    instruct=None,
+    read_text_only=True,
+    prompt_speech=None,
+    audio_list=None,
+):
+    res = get_tts_sft_prompt(
+        text,
+        instruct=instruct,
+        read_text_only=read_text_only,
+        prompt_speech=prompt_speech,
+    )
+
+    prompt = to_prompt(res)
+    final_prompt = {
+        "prompt": prompt,
+    }
+    if audio_list is not None:
+        final_prompt.update(
+            {
+                "multi_modal_data": {
+                    "audio": audio_list,
+                },
+            }
+        )
+    return final_prompt
+
+
+def get_audio_understanding_sft(audio_path, text="", thinking=False, use_sostm=False):
+    audio_list = []
+    audio_list.append(get_audio_data(audio_path))
+    res = get_audio_understanding_sft_prompt(
+        input_speech=audio_path, input_text=text, thinking=thinking, use_sostm=use_sostm
+    )
+    prompt = to_prompt(res)
+    final_prompt = {
+        "prompt": prompt,
+        "multi_modal_data": {
+            "audio": audio_list,
+        },
+    }
+    return final_prompt
+
+
+def get_spoken_dialogue_sft_multiturn(message_list, system_prompt=None, ref_audio_path=None, audio_list=None):
+    res = get_spoken_dialogue_sft_multiturn_prompt(
+        message_list, system_prompt=system_prompt, prompt_speech=ref_audio_path
+    )
+    prompt = to_prompt(res)
+    final_prompt = {
+        "prompt": prompt,
+        "multi_modal_data": {
+            "audio": audio_list,
+        },
+    }
+    return final_prompt
+
+
+def get_speech2text_dialogue_sft_multiturn(message_list, thinking=False, audio_list=None):
+    res = get_s2t_dialogue_sft_multiturn_prompt(
+        message_list,
+        thinking=thinking,
+    )
+    prompt = to_prompt(res)
+    final_prompt = {
+        "prompt": prompt,
+        "multi_modal_data": {
+            "audio": audio_list,
+        },
+    }
+    return final_prompt
+
+
+def get_text_dialogue_sft_multiturn(
+    message_list,
+):
+    res = get_text_dialogue_sft_multiturn_prompt(
+        message_list,
+    )
+    prompt = to_prompt(res)
+    final_prompt = {
+        "prompt": prompt,
+    }
+    return final_prompt
+
+
+query_map = {
+    "tts_sft": get_tts_sft,
+    "tts_sft_with_instruct": get_tts_sft,
+    "tts_sft_with_audio": get_tts_sft,
+    "tts_sft_with_natural_instruction": get_tts_sft,
+    "audio_trancribing_sft": get_audio_understanding_sft,
+    "audio_understanding_sft": get_audio_understanding_sft,
+    "audio_understanding_sft_with_thinking": get_audio_understanding_sft,
+    "spoken_dialogue_sft_multiturn": get_spoken_dialogue_sft_multiturn,
+    "speech2text_dialogue_sft_multiturn": get_speech2text_dialogue_sft_multiturn,
+    "text_dialogue_sft_multiturn": get_text_dialogue_sft_multiturn,
+}
+
+
+def main(args):
+    model_name = args.model_name
+
+    # Get paths from args
+    text = getattr(args, "text", None)
+    audio_path = getattr(args, "audio_path", None)
+
+    instruct = getattr(args, "instruct", None)
+
+    # Get the query function and call it with appropriate parameters
+    query_func = query_map[args.query_type]
+
+    omni_llm = Omni(
+        model=model_name,
+        stage_configs_path=args.stage_configs_path,
+        log_stats=args.enable_stats,
+        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
+        init_sleep_seconds=args.init_sleep_seconds,
+        batch_timeout=args.batch_timeout,
+        init_timeout=args.init_timeout,
+        shm_threshold_bytes=args.shm_threshold_bytes,
+    )
+
+    thinker_sampling_params = SamplingParams(
+        temperature=0.0,
+        top_p=1.0,
+        top_k=-1,
+        max_tokens=2048,
+        seed=SEED,
+        logit_bias={},
+        repetition_penalty=1.1,
+    )
+
+    code2wav_sampling_params = SamplingParams(
+        temperature=0.0,
+        top_p=1.0,
+        top_k=-1,
+        max_tokens=4096 * 16,
+        seed=SEED,
+        detokenize=True,
+        repetition_penalty=1.1,
+    )
+
+    sampling_params_list = [
+        thinker_sampling_params,
+        code2wav_sampling_params,
+    ]
+
+    # Build query result based on query type
+    # Notice: The audio files used in this example are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples
+    if args.query_type == "tts_sft":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type tts_sft
+        query_result = query_func(text=text, read_text_only=True)
+    elif args.query_type == "tts_sft_with_instruct":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type tts_sft_with_instruct --instruct "Speak happily in a child's voice"
+        query_result = query_func(text=text, instruct=instruct, read_text_only=True)
+    elif args.query_type == "tts_sft_with_audio":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type tts_sft_with_audio --audio_path "./spoken_dialogue_assistant_turn_1.wav"
+        audio_list = [get_audio_data(audio_path)]
+        query_result = query_func(text=text, read_text_only=True, prompt_speech=audio_path, audio_list=audio_list)
+    elif args.query_type == "tts_sft_with_natural_instruction":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type tts_sft_with_natural_instruction --text "In a panting young male voice, he said: I can't run anymore, wait for me!"
+        query_result = query_func(text=text, read_text_only=False)
+    elif args.query_type == "audio_trancribing_sft":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type audio_trancribing_sft --audio_path "./spoken_dialogue_assistant_turn_1.wav"
+        audio_path = "spoken_dialogue_assistant_turn_1.wav"
+        text = "Please transcribe this audio and repeat it once."
+        query_result = query_func(text=text, audio_path=audio_path, use_sostm=True)
+    elif args.query_type == "audio_understanding_sft":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type audio_understanding_sft --text "Summarize the audio." --audio_path "./spoken_dialogue_assistant_turn_1.wav"
+        query_result = query_func(text=text, audio_path=audio_path)
+    elif args.query_type == "audio_understanding_sft_with_thinking":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type audio_understanding_sft_with_thinking --text "Summarize the audio." --audio_path "./spoken_dialogue_assistant_turn_1.wav"
+        query_result = query_func(text=text, audio_path=audio_path, thinking=True)
+    elif args.query_type == "spoken_dialogue_sft_multiturn":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type spoken_dialogue_sft_multiturn  --audio_path "./prompt_speech_zh_m.wav"
+        first_turn_text_response = "我没办法获取实时的天气信息。不过呢，你可以试试几个方法来查看今天的天气。首先，你可以用手机自带的天气功能，比如苹果手机的天气应用，或者直接在系统设置里查看。其次，你也可以用一些专业的天气服务，像是国外的AccuWeather、Weather.com，或者国内的中国天气网、墨迹天气等等。再有就是，你还可以在谷歌或者百度里直接搜索你所在的城市加上天气这两个字。如果你能告诉我你所在的城市，我也可以帮你分析一下历史天气趋势，不过最新的数据还是需要你通过官方渠道去获取哦。"
+        audio_list = []
+        s1_audio_path = "weather_of_today.mp3"
+        s2_audio_path = "spoken_dialogue_assistant_turn_1.wav"
+        s3_audio_path = "beijing.mp3"
+        audio_list.append(get_audio_data(audio_path))
+        audio_list.append(get_audio_data(s1_audio_path))
+        audio_list.append(get_audio_data(s2_audio_path))
+        audio_list.append(get_audio_data(s3_audio_path))
+
+        message_list = [
+            {"role": "user", "content": s1_audio_path},
+            {"role": "assistant", "content": {"text": first_turn_text_response, "audio": s2_audio_path}},
+            {"role": "user", "content": s3_audio_path},
+        ]
+        query_result = query_func(message_list, system_prompt=None, ref_audio_path=audio_path, audio_list=audio_list)
+    elif args.query_type == "speech2text_dialogue_sft_multiturn":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type speech2text_dialogue_sft_multiturn
+        s1_audio_path = "weather_of_today.mp3"
+        s2_audio_path = "beijing.mp3"
+        audio_list = []
+        audio_list.append(get_audio_data(s1_audio_path))
+        audio_list.append(get_audio_data(s2_audio_path))
+        message_list = [
+            {"role": "user", "content": s1_audio_path},
+            {
+                "role": "assistant",
+                "content": "Hello, I can't get real-time weather information. If you can tell me your city, I can help you analyze historical weather trends, but for the latest data, you'll need to get it through official channels.",
+            },
+            {"role": "user", "content": s2_audio_path},
+        ]
+        query_result = query_func(message_list, thinking=True, audio_list=audio_list)
+    elif args.query_type == "text_dialogue_sft_multiturn":
+        # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH}  --query-type text_dialogue_sft_multiturn
+        message_list = [
+            {"role": "user", "content": "Could you recommend some tourist attractions in China?"},
+            {"role": "assistant", "content": "Hello, which city would you like to travel to?"},
+            {"role": "user", "content": "Beijing"},
+        ]
+        query_result = query_func(message_list=message_list)
+    else:
+        raise ValueError(f"Invalid query type: {args.query_type}")
+
+    prompts = [copy.deepcopy(query_result) for _ in range(args.num_prompts)]
+
+    print("prompts", prompts)
+    omni_outputs = omni_llm.generate(prompts, sampling_params_list)
+
+    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
+    if args.query_type is not None:
+        output_dir = os.path.join(output_dir, args.query_type)
+    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 = output.request_id
+                text_output = output.outputs[0].text
+                # Save aligned text file per request
+                prompt_text = output.prompt
+                out_txt = os.path.join(output_dir, f"{request_id}.txt")
+                lines = []
+                lines.append("Prompt:\n")
+                lines.append(str(prompt_text) + "\n")
+                lines.append("vllm_text_output:\n")
+                lines.append(str(text_output).strip() + "\n")
+                try:
+                    with open(out_txt, "w", encoding="utf-8") as f:
+                        print("lines", lines)
+                        f.writelines(lines)
+                except Exception as e:
+                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
+                print(f"Request ID: {request_id}, Text saved to {out_txt}\n")
+        elif stage_outputs.final_output_type == "audio":
+            for output in stage_outputs.request_output:
+                request_id = output.request_id
+                audio_tensor = output.outputs[0].multimodal_output.get("audio")
+
+                if audio_tensor is None:
+                    continue
+
+                output_wav = os.path.join(output_dir, f"{request_id}.wav")
+
+                # Convert to numpy array and ensure correct format
+                audio_numpy = audio_tensor.float().detach().cpu().numpy()
+
+                # Ensure audio is 1D (flatten if needed)
+                if audio_numpy.ndim > 1:
+                    audio_numpy = audio_numpy.flatten()
+
+                # Save audio file with explicit WAV format
+                sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
+                print(f"Request ID: {request_id}, Audio saved to {output_wav}")
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser.add_argument(
+        "--model-name",
+        "-m",
+        type=str,
+        default="XiaomiMiMo/MiMo-Audio-7B-Instruct",
+        help="Backbone LLM path.",
+    )
+    parser.add_argument(
+        "--text",
+        "-t",
+        type=str,
+        default="The weather is so nice today.",
+        help="input text",
+    )
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="tts_sft",
+        choices=query_map.keys(),
+        help="Query type.",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to local audio file. If not provided, uses default audio asset.",
+    )
+    parser.add_argument(
+        "--instruct",
+        type=str,
+        default=None,
+        help="instruct",
+    )
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        help="Enable writing detailed statistics (default: disabled)",
+    )
+    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(
+        "--batch-timeout",
+        type=int,
+        default=5,
+        help="Timeout for batching in seconds (default: 5)",
+    )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=5000,
+        help="Timeout for initializing stages in seconds (default: 5000)",
+    )
+    parser.add_argument(
+        "--shm-threshold-bytes",
+        type=int,
+        default=65536,
+        help="Threshold for using shared memory in bytes (default: 65536)",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default="./output_audio",
+        help="Output audio wav directory.",
+    )
+    parser.add_argument(
+        "--output-wav",
+        default="output_audio",
+        help="[Deprecated] Output wav directory (use --output-dir).",
+    )
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=1,
+        help="Number of prompts to generate.",
+    )
+
+    parser.add_argument(
+        "--sampling-rate",
+        type=int,
+        default=24000,
+        help="Sampling rate for audio.",
+    )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default="../../../model_executor/stage_configs/mimo_audio.yaml",
+        help="Path to a stage configs file.",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

examples/offline_inference/mimo_audio/message_base64_wav.json

@@ -0,0 +1,56 @@
+{
+  "messages": [
+    {
+      "role": "system",
+      "content": [
+        {
+          "type": "text",
+          "text": "Your voice should be:"
+        },
+        {
+          "type": "audio_url",
+          "audio_url": {
+            "url": "data:audio/wav;base64,UklGRiRKAwBXQVZFZm10IBAAAAABAAEAwF0AAIC7AAACABAAZGF0YQBKAwAFAP3/CAD7/wkAAAD9/wgA+v8JAP//AQADAP//CAD+/wEABAABAP7/BgD8//7/BgD0/wYA/f/8/wMA9f8GAP///P8DAPz/AAAFAP3//v8CAP7////+////+/8IAP//AQAAAAMABgD6/wUA+v8DAAcA+/8AAAMAAAAAAP3//v8DAAMACgD8/wQACQD9/wUA/f8BAAcABQADAAAABAADAAEA+v8AAAIAAAAFAPz//f8DAAAA/f8AAAMA/v8FAP7/9v8GAP7//f8EAAAAAQAIAPv/+/8AAP3////+//3/AAAHAPr//v8BAPz/AwD///n/AAACAPr//f/8//7/AAD//wAA+f8GAPz/9v8CAPX/BgACAPv/AQD//wQA///7/wAAAAAHAAQA/f8FAAAAAgAAAPv/BAAGAAAABgACAP7/DQD7//3/BAD//wkA//8DAPz/AwAHAPf/AwAAAAAABAD7//f/AwD8/wAA/f/5/wQA/P8AAPr//P8AAAAA//8AAP////8DAPL/AQD9//r/CwD4/wEAAgD9//7//P/9//7/AAD//wAA/v8DAPr/+v8CAPr///8CAPz///8GAPr/AAAEAP//AgADAP//AAAHAPz/BgD//wUACwD//wYAAgABAAQABQD8/wgABgACAAUA//8CAAQAAAADAP7/AQAIAPr/AwD+//v/AgACAPv/AQAJAPf/AQD+//n/AgAAAP///f8CAAAA/P8DAPf/+/8HAPv///////7//f8AAP3//P8BAPr/AAD6//z/AAD5/wIAAAD7/wIAAgD4/wAA/v/6/wQAAQAAAAAABwAAAAAAAwD+/wIABQADAAEABAADAAAAAwABAAAABAAGAAEAAgAGAAAABQD///7/BQAAAAMAAQAAAAIABAD8/wMA/////wQA/f8AAAEAAQD//wMA+/8FAPr///////P/CwDz/wQA/P/6/wIA8/8AAPH//P/+//f//v/+//n//P/8//L//v/7//r////7//3//v/5//v/AAD5/wMAAAD8/wYA/P///////////wAABwD6/wwAAQD+/wUA/P8GAP7/AwD+/wUABwD+/wUA//8FAAMAAgD//wUAAAAEAAcA+/8KAAUAAgAEAAAA//8HAAAA+v8EAAIAAQAEAPz/+/8EAP7//////wAAAQABAP//+///////AAD9/wEA/f8DAAAA+P8AAPj/AQD+//v/AQD8/wIA/v/6/wAA/////wIA/P///wMA+/8BAPj///8CAAAAAwD+/wcAAAAFAP3/AAAFAP7/AwABAP7/AAACAPv/AwD+/wMABAD9/wUA/v8AAAEA/v8AAAIAAQAAAAAAAAD8/wIA/v/8/wMAAAD///3/AAD5//7/AwD2/wgAAAD8/wQA+f8CAP7/AAAAAP//AwD//wAA/v//////AAAAAAAAAwD9/wIA+//9/wQA+/8DAAAA//8AAP///f/+/wAAAAABAP//BgD8/wIAAgD3/wwA//8AAAkA/f8EAAAA//8BAAAAAwADAP//BQD///z/AwD8/wEAAwAAAAEAAQAAAAAA/v8BAAAAAQAGAP7/AAD///3//v8AAAAAAgAGAP//AgAAAP3////9/wAAAQACAAAAAQD///3/AQD9/wAABwD//wMABAD8/wEAAwD+/wIABQD+/wEAAwD6////BAD5/wYAAAD+/wgA+f8CAP3///8DAP7/AgAAAAAA/v8EAPv/AAAHAPv/CQACAP7/BgD//wIAAAADAAAAAgAEAPz/BQD8/wQA/P/+/wMA+f8EAP3//P8AAP////8AAPv/AgD8/wAAAgDz/wgA/f/7/wQA+P8BAAAA+v8BAP7/AAACAP3/AAACAP7/AQAEAPj/BwAAAAAABgD7/wcA+v8CAP7/+/8IAP3/AQACAAAAAQACAAAA/v8CAAcA/f8HAAIA/v8IAPv/AQAAAAAABAD+/wQAAgD//wMA/P8AAAAA//8BAPr/BQD9/wAA///4/wIA+/8AAP3//v8EAPr/AAD9//r/BQD3/wEA+v/6/wYA8/8BAPz/9v8CAPn/+f/+//3//v/+//3//v///wAA/P/7/wAAAAD//wAA//8AAP7/AgD4////BAD7/wYAAAADAAAAAwAAAPr/BQD6/wUABAD//wcA/f8EAPz//P8FAPv/BQADAP//BgAAAP//AQAAAAAAAgAAAAEAAwAAAAAA//8AAAAAAAAAAP7/BAD///7/AgD8/wAAAQD8////AQD9/wQA/v///wIA//8BAPv/AQD9/wMA/v/8/wQA+f8DAP3/9/8EAAAA/P8GAPz/AwAEAP//AAD8/wQAAAAGAAAAAgAFAAEABAD8/wEAAQADAAAAAgADAAAACAD7//7/BQD9/wUAAAD+/wcAAAABAAAA/f8DAAMAAAAAAAEA/v8AAAAA+f8AAAEA//8AAP//AAD7/wUA9//6/wQA9/8BAP3//f8AAP7/AAD7//3/AwD+//v//v/7//3/BAD3//r/CAD8/wIA/v/6//3/AQD+//j/AAAAAAgAAAD///v//f8BAP7//v/9/wEA//8EAAAA//8BAAEAAAD+/////f8AAP7/AgD9/wQABQD4/wIA/f/5/wAA/P/9/wEAAQABAPf/AAD8//7/AQD5/wAA/f8BAPf/+f/7//b/AwADAPz///8BAPn//P/4/////P8AAAAA+P8JAAAA/f/8//3/AAABAAAA8v8EAAMAAAACAPn/BAAOAAcA+//t////DwAEAAIA/f8HAAYADAD7//f/BgD//wwA//8BAAEADwAEAPv/BQD1/xEA//8DAP//9/8LAPn/AgABAAMABAAEAPz//f8FAPb/BAD+//r/CwAAAP7/BQD4////BAD4/wEA/v8AAAAA/v/+//z/BAABAP//BQD+//z/AQD7/wAA/P8DAAUAAAAHAP//AQD9////AAABAP//AAACAAQACAD+/wYA//8DAAcA/v8CAAEABQADAAIAAgACAAUABQAAAAAAAgAAAAAAAwABAAAABQAAAAYABgAAAAEA/v8EAAAA/P8GAP7/BAAGAPn/AQABAAIA/f8EAPn///8GAPT/BgD5/wAAAgD8/wQA+v8BAPv/9/8AAP3/AAAAAPz/AQD9//r////6////AAD9/////f8AAPz/AAD9//z/CQD5/wUA+v/7/wIA9P8FAPn/AQADAP3//v8AAP3/+/8DAPr/AQABAAAA/f8BAAAA9/8EAP7///8GAPr///8BAPz/AQD9/wAABQD//wEAAAD9/wMAAQD7/wAAAwAAAAMAAAD+/wMAAQAAAP3/AAABAP//AgD+/wEAAQAHAPv///8EAPb/CQD//wAAAAAAAAUA9v8JAPf///8FAPr/AwD9/wIA+v8DAP//+P8GAP////8EAPv/+/8EAPv/AAD+//7/AQD//wEA+/8AAAAA/v8CAP///P8GAP////8DAPv/BgAAAAAAAgD//wEA//8CAP7/AgAAAP//BgAAAAEAAgAAAP//BwD6/wIABAD6/woA+/8EAAAAAAAEAPv/BgD//wIAAAABAAMAAAAHAPf/BwAAAPr/BwD4/wQAAAAAAAAAAAD//wAAAAD7/wcA/v8AAAAA///+/wIA///7/wkA+/8GAAAA/f8BAP//AgD9/wYA//8FAAQA//8CAPz/AQAFAPr/AwAAAP//BgD5//7//v8CAPv/AAD9//b/CAD4//z//v/6/wAAAAD+//z/AAABAP///v////3/BwABAAAAAAAAAAUA/f8DAPb/AAAIAPv/AwD9////BAAAAPf/+/8AAPv/AQD7//v/AAABAP7/+v8BAPn/AQACAPr/AQACAAEA/f8FAPv/AAAFAPj/AAABAPv//P/+//n//v////r//f/+//r/AAD0//b//f/0/wMA9//+//7//f8BAPP/BAD6////AwD8/wEAAAAKAPf/BwACAPz/DAD9//3/AwAGAPn/BwD+//3/DgD7/wAAAAD7/wYA/f/5/wQA+v8GAAgA+P8IAAMAAwAKAP//AQAFAAMABQAAAAUABQAGAAcA//8IAPr/AgAAAPX/BAD+//3/BAAAAPX/AwD5//T/AwD3//n/BgD6//z/BAD3/wAAAAD+/wAA//8AAAAA/v/+//7///8BAAAAAAABAP7//v/+//b/AAD8/wAAAAD+/wEA/f8CAPr/AQD7/wAABQD7/wcA//8FAAcAAgAIAAUABAAGAAcAAwAGAAQABAAKAAkABQAIAAUABwADAAMAAAACAAMAAgAFAP7/DAABAAgA/v/+/wYA/P8KAPj/BAAGAP7/CQAAAAEACAAEAAUA/P///wIA+/8DAP7///8FAP//AAD7/wAA/f/2/wIA9P/9/wMA8v///wAA+/8DAAEA/P8CAP7//v/5//r/CAAAAAkABgAAAAUABwD+//3/CwD//w4ABwAAAAkAAgAJAPz/BwALAAEAEgD//wAACgABAAcAAgAFAAcABwAGAAAABAAEAAoAAwAIAAsABQAMAAAAAQACAAkABwAGAA4ACwAKABMAAgAAAAoA/f8JAAAABAADAAsACQD9/wQA+/8AAPf//v/z//r/AQDv//3/8f/w//r/8f/0/+//9P/u/+j/8v/g/+7/8f/r//X/8f/x/+7/7v/m/+v/7v/3//D/8//4/+3//v/z//b/9//7//n/9f8BAPT//P8AAPj/AAADAAAAAAD+/wIA+f/8/wQA+v8DAAcA/v8AAA8AAAAFABEA//8NAA4ABwAHABAAEgAKABgAFAARABYAFAARABIAHAAXABIAIAAYABQAIgAUABEAIQAXABUAGgAVABMAFwAVAAgADQAJAAkABAAJAAkAAAAKAPz/+f/3//z/8P/2/wAA8/8CAP7/8//0//j/7v/t/Line truncated
+          }
+        }
+      ]
+    },
+    {
+      "role": "user",
+      "content": [
+        {
+          "type": "audio_url",
+          "audio_url": {
+            "url": "data:audio/wav;base64,UklGRgrVAABXQVZFZm10IBAAAAABAAEAwF0AAIC7AAACABAAZGF0YebUAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA//8AAAAAAAAAAAAAAAAAAAAAAAAAAP////8AAAAAAAAAAAAAAAD//wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP////8AAAAAAAAAAAAA///p/9T/xv/G/8b/yf/K/8r/zP/M/83/zf/O/83/zf/P/8//z//O/87/z//P/8//z//P/8//z//P/8//z//P/8//z//Q/9D/0P/Q/9D/0P/R/9H/0f/R/9H/0f/R/9H/0f/R/9L/0v/S/9L/0v/S/9P/0//T/9P/0//T/9P/0//U/9T/1P/U/9T/1P/V/9T/1f/V/9X/1f/W/9b/1v/W/9b/1v/W/9b/1//X/9f/1//Y/9j/2P/Y/9j/2P/Z/9n/2f/Z/9n/2f/a/9r/2v/a/9v/2//b/9v/2//b/9v/3P/c/9z/3f/c/93/3f/d/93/3f/e/97/3v/e/9//3//f/9//4P/g/+D/4P/g/+D/4P/h/+H/4f/h/+H/4f/i/+L/4v/j/+P/4//j/+P/4//k/+P/5P/k/+T/5P/l/+X/5f/l/+X/5f/l/+X/5f/l/+X/5v/n/+f/5//o/+j/6P/o/+j/5v/n/+f/5//o/+j/6P/o/+n/6f/p/+j/6P/o/+j/6P/p/+n/6f/p/+n/6v/q/+r/6v/q/+r/6v/q/+r/6v/q/+r/6//r/+r/6v/r/+r/6v/p/+r/6v/r/+r/6v/r/+z/7P/s/+v/7P/r/+v/6//s/+v/6//r/+v/6//t/+z/7f/t/+3/7f/r/+z/6//s/+v/7P/s/+z/7P/s/+z/7P/t/+3/7f/s/+z/6//s/+r/6//q/+z/7P/t/+3/7v/u/+3/7f/t/+z/7P/s/+z/7P/t/+3/7f/t/+3/7f/t/+3/7f/t/+z/7P/s/+z/7f/t/+3/7f/u/+3/7f/t/+3/7f/t/+3/7f/u/+3/7f/t/+3/7f/t/+3/7f/t/+3/7P/r/+z/6//t/+z/7f/t/+7/7f/u/+7/7f/t/+z/7P/t/+z/7P/t/+3/7P/t/+7/7v/v/+7/7v/t/+z/7P/t/+3/7v/u/+7/7v/v/+//7v/u/+3/7P/u/+z/7P/s/+z/7f/u/+7/7v/v/+7/7v/s/+z/7P/t/+3/7f/u/+7/7f/u/+3/7v/t/+3/7f/t/+3/7P/t/+z/7f/u/+3/7//u/+7/7f/t/+3/7f/u/+//7v/v/+7/7v/t/+7/7f/v/+7/7v/t/+7/7f/t/+3/7v/u/+7/7v/v/+3/7v/t/+z/7P/r/+z/6//t/+r/7P/t/+7/7v/u/+7/7f/s/+z/6//s/+v/7P/t/+3/7//w//D/7v/w/+7/7v/s/+3/6//t/+v/7f/u/+//7//u/+//7v/u/+r/6//q/+v/6//s/+7/7v/w/+//7//v/+3/7f/s/+v/7P/s/+3/7f/u/+7/7v/u/+3/7P/r/+z/7f/s/+z/7P/t/+3/7P/t/+3/7v/s/+3/7f/t/+z/6//r/+z/7P/s/+z/7P/t/+z/7f/t/+z/7v/v/+7/7P/s/+z/6//u/+3/7v/u/+7/7f/u/+3/7f/t/+v/7f/r/+v/6//s/+3/7P/u/+7/7v/v/+7/7v/u/+z/7P/r/+3/6//t/+3/7v/t/+7/7v/v/+7/7v/t/+3/7P/t/+7/7v/v/+3/7v/t/+3/7f/s/+z/7f/s/+3/7f/t/+7/7v/u/+3/7//u/+7/7f/s/+z/7P/s/+v/7P/t/+3/7f/v/+7/7v/u/+3/7f/t/+3/7f/u/+7/7v/u/+7/7//u/+7/7v/t/+3/7f/t/+3/7f/u/+//7//v/+//7v/u/+7/7v/u/+7/7v/u/+7/7v/t/+7/7v/u/+7/7v/u/+7/7f/t/+3/7f/t/+7/7f/t/+3/7P/u/+3/7v/t/+7/7v/u/+7/7v/v/+7/7v/v/+7/7v/u/+7/7v/u/+3/7f/u/+3/7v/u/+7/7f/u/+3/7v/u/+7/7f/t/+7/7v/v/+7/7v/u/+3/7v/t/+3/7f/t/+3/7f/u/+7/7v/u/+3/7v/t/+3/7P/t/+3/7f/u/+3/7v/u/+//7v/u/+7/7v/u/+7/7v/u/+7/7f/u/+3/7f/u/+7/7f/u/+3/7f/t/+3/7f/t/+3/7f/s/+z/7f/t/+3/7f/t/+3/7P/t/+3/7f/t/+z/7P/s/+3/7P/t/+3/7f/u/+3/7f/t/+3/7f/t/+7/7P/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/s/+z/7f/t/+3/7f/t/+3/7f/t/+z/7P/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/t/+z/7P/s/+3/7f/t/+3/7f/t/+3/7P/s/+z/7P/t/+z/7P/s/+z/7P/s/+3/7P/t/+3/7f/t/+3/7f/t/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/t/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+v/6//s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7f/s/+z/7P/s/+z/7P/s/+v/6//r/+z/6//s/+z/7P/s/+z/7P/s/+z/7P/s/+z/6//s/+z/7P/s/+v/6//s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+3/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/s/+z/7f/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/s/+3/7P/t/+3/7P/s/+z/7P/t/+z/7f/t/+3/7f/t/+3/7f/u/+z/7f/t/+3/7f/t/+3/7f/u/+3/7f/t/+z/7P/s/+z/7P/t/+z/7v/t/+7/7f/t/+7/7f/v/+z/7f/r/+3/7P/t/+3/7f/u/+7/7v/v/+7/7v/u/+z/7f/s/+7/7f/u/+3/7f/v/+7/7//u/+7/7//t/+//7v/u/+z/7P/s/+7/7f/w/+3/7f/u/+//7f/u/+7/7//s/+//6//w/+3/7//u/+3/7v/s/+7/7P/u/+//7v/x/+z/7//t/+3/7f/u/+//7//u/+7/6//w/+3/8f/t/+//7v/v/+7/7v/u/+//7v/w/+3/7v/t/+//7//t/+7/7v/w/+3/7P/v/+3/7//t/+//7v/w/+7/7//y/+3/7v/v/+3/7//t/+7/7P/v/+3/9P/w/+3/6//r//D/7//w/+3/8v/t/+7/7f/v/+3/7f/v//D/7f/t//D/7v/u/+3/8P/v/+7/7v/w/+z/7//r//L/7v/x/+3/7v/u/+7/7//t/+//7f/v/+7/7f/v/+7/7f/u/+7/7f/v/+//8P/u/+3/7//u/+//7f/v/+z/7f/v/+3/7v/p//D/7f/x/+7/7v/s/+3/8P/v/+//7f/r/+7/7P/t/+3/7v/t/+3/7f/u/+7/7//t/+7/7f/u/+3/6//s/+z/7f/t/+3/7//s/+//7f/y/+z/7v/s/+3/7P/v/+3/6//r/+3/7P/u/+//7f/t/+7/7v/v/+z/7f/r/+//7f/w/+z/7f/v/+v/7f/s/+7/7v/v/+z/7P/u/+z/7//s/+3/7f/u/+7/7f/w/+z/7v/t/+3/7v/t/+3/7P/u/+7/7v/s/+z/7f/v/+3/7v/t/+3/7f/t/+7/7P/t/+3/7v/v/+z/7f/u/+3/7v/v/+z/7v/t/+//7f/u/Line truncated
+          }
+        }
+      ]
+    },
+    {
+      "role": "assistant",
+      "content": [
+        {
+          "type": "text",
+          "text": "我没办法获取实时的天气信息。不过呢，你可以试试几个方法来查看今天的天气。首先，你可以用手机自带的天气功能，比如苹果手机的天气应用，或者直接在系统设置里查看。其次，你也可以用一些专业的天气服务，像是国外的AccuWeather、Weather.com，或者国内的中国天气网、墨迹天气等等。再有就是，你还可以在谷歌或者百度里直接搜索你所在的城市加上天气这两个字。如果你能告诉我你所在的城市，我也可以帮你分析一下历史天气趋势，不过最新的数据还是需要你通过官方渠道去获取哦。"
+        },
+        {
+          "type": "audio_url",
+          "audio_url": {
+            "url": "data:audio/wav;base64,UklGRiQiKQBXQVZFZm10IBAAAAABAAEAwF0AAIC7AAACABAAZGF0YQAiKQAGAAcABwAIAAkACQAJAAoACwAMAAwADQAOAA4ADwAPABAAEAAQABAAEQAQABAAEQARABAAEQARABEAEAAQABEAEQAQABAAEAAPAA4ADwAOAA0ADQANAAwACwAJAAYABQAEAAIAAQAAAP7//f/7//r/+v/5//b/9v/0//X/9f/0//X/9P/0//L/9P/y//D/9P/z//T/8v/x//L/8v/x//H/8P/v/+//7v/u/+3/7P/t/+3/7P/t/+7/7v/u//D/8f/z//X/9f/3//n/+//7//3///8AAAEABAAGAAgACgAMAA4AEAASABMAFAATABQAFAAUABUAFQATABIAEgARABEAEAAOAAwACwAJAAcABgAEAAMAAgABAAEAAAAAAAAA//////7//f/8//v/+f/4//j/9//3//f/9v/2//b/9f/1//b/9f/1//b/9v/2//b/9f/1//X/9P/z//P/9P/z//P/8//0//T/9P/0//X/9v/2//f/+P/4//n/+v/7//z//f/+//7///8AAAEAAgACAAMABAAFAAYABwAJAAoADAANAA4ADwAPAA8AEAAQABEAEgASABMAFAAVABYAFwAXABYAFgAWABYAFgAVABQAEwATABMAEgASABIAEQARABAADwAOAA4ADQANAAsACgAIAAcABAACAAEA///9//v/+f/4//b/9f/z//H/8f/v/+3/7P/r/+n/5//m/+X/5f/k/+T/5P/l/+X/5//o/+r/7P/u//D/8v/1//f/+f/7//z//v8AAAEAAgAEAAQABQAHAAUABQAGAAQAAwACAAEAAgAAAAEAAQD+/wAA/v8AAP7//P////7//f/9/wAA/v/9//7//v///wAA//////////8AAAEAAwACAAIAAwADAAMAAgACAAIAAwACAAEAAAABAAEAAAD////////////////+/////v/9//3//P/8//z//f/9//3//v/+//7///8AAP////////////////////////7//v/9//3//f/8//r/+P/3//j/9//3//b/9f/2//b/9f/1//b/9//3//f/9//3//f/9//3//f/9//3//f/9//3//f/9//3//f/9//3//f/9//4//j/+P/5//n/+f/5//n/+f/5//n/+f/6//r/+//7//z//P/+////AAACAAMABQAGAAcACAAIAAkACQAJAAkACQAKAAoACQAJAAcABgAEAAMAAQD///3/+//6//n/+P/3//b/9v/1//b/9v/2//b/9v/2//b/9//4//n/+v/8//7//v///wEAAwAEAAUABwAJAAoACgALAA4ADAANAA0ADwAOAA0ADwAPAA8ADgARABEAEQARABEAEAAQABEAEQARABAAEAAQABEAEQAQABAADwAPAA4ADQAOAAwADQANAAwADQANAAsADAALAAsACgAJAAkACAAHAAcABQAFAAUAAgABAAEAAAD///7//P/6//n/9v/z//P/8v/w/+//7f/r/+v/6v/p/+r/6f/o/+j/6P/o/+j/6f/p/+n/6v/r/+z/7P/t/+7/7//w//H/8//1//b/+P/5//r//P/9////AAABAAIAAwAFAAYABwAIAAkACwALAAwADQAOAA4ADwAQABEAEgASABMAFAAVABUAFQAVABYAFQAVABUAFAATABIAEQAPAA4ADQAMAAsACgAJAAcABQADAAAA///8//r/+P/2//X/8//x/+//7v/t/+z/6//r/+r/6v/q/+v/6//s/+z/7f/t/+7/7//w//H/8v/z//T/9f/2//f/9//4//j/+f/6//r/+//7//z//P/9//7//v///wAAAAAAAAEAAQACAAMABAAGAAcACQAKAAwADQAPABAAEQASABIAEgATABUAFwAYABkAGgAbAB0AHgAfAB4AHgAdABwAGwAaABkAGAAWABQAEgAQAA0ACwAJAAcABQADAAEA/v/8//r/+P/2//T/8//x//D/7//t/+z/6//r/+v/6//r/+v/6//s/+3/7f/t/+7/7v/v/+//8P/w//D/8f/w//D/7//w//H/8f/w//D/8f/y//L/8v/y//P/8//1//f/+P/5//z//f/8/wAA//8DAAEAAAAEAAIAAQADAAMAAQAAAAEAAAABAAEA/////////v/9//7////9//z//P/9//3//P/8//z//f/+//7//v////////8AAAAAAQACAAIAAgACAAIAAgABAAEAAQABAAIAAQAAAAAAAAD///7//v/9//3/+//6//r/+v/6//r/+f/6//r/+v/5//n/+v/7//v/+v/6//v//f/+//7//v//////AAD///////////7//f/8//v/+//6//r/+f/5//n/+v/6//r/+v/6//z//P/9//3///8AAAEAAgADAAUABQAGAAYABwAIAAkACQAKAAsADAAMAAwADAAMAA0ADAAMAAsACwAKAAkACAAHAAUAAwACAAEAAAD+//3/+//6//j/9//1//T/8//y//L/8//z//T/9f/1//b/9v/2//f/+P/4//n/+f/5//r//P/9////AAACAAMABAAFAAYACAAJAAoACwANAA4ADQAOABEAEAARABEAFAATABMAFQAVABYAFQAXABYAFgAWABUAFAATABQAEwASABIAEQAQABAADwAPAA4ADQAMAAsACwAMAAoACwAKAAkACgAJAAgABwAFAAUAAwACAAEA///9//3/+v/5//j/9P/y//D/7v/s/+v/6P/m/+X/4v/h/+H/4f/g/+H/4P/g/+H/4v/j/+T/5f/m/+f/6P/q/+v/7f/u/+//8f/y//T/9f/3//j/+v/7//z//v///wEAAgADAAQABQAGAAgACQAKAAsADAAOAA8AEAASABMAFAAVABYAFgAXABgAGAAZABkAGgAaABoAGQAZABgAFgAVABMAEgAQAA8ADgAMAAoACQAHAAUABAACAAAA///9//v/+f/3//X/8//x//D/7//v/+7/7v/t/+3/7f/t/+7/7v/v//D/8P/x//L/8//0//X/9f/2//f/9//4//j/+f/5//r/+v/6//r/+v/6//r/+f/6//r/+v/7//v//P/8//3//v///wEAAgADAAUABgAIAAsADQAPABEAEwAUABYAFwAYABgAGQAZABkAGAAZABkAGQAZABgAFwAVABQAEgARAA4ACwAIAAUAAwABAP///v/8//v/+v/4//b/9f/0//P/8v/y//H/8P/w/+//7v/u/+3/7f/t/+3/7f/s/+z/7P/t/+3/7v/u/+7/7v/u/+//8P/v//D/8P/x//H/8f/y//P/9P/0//X/9f/1//f/+P/5//n/+//8//z//f/9//7//f///wAAAAAAAAEAAQD//wAA/f8AAPz/+//+//v/+v/8//z/+//7//z//f/+/////v///wEAAQAAAAIABAADAAIAAwADAAQAAwADAAMAAwADAAIAAgACAAEAAAAAAAAAAAD//////v/+//7//v/9//7//v/+//////////////////////8AAAAA////////AAAAAAAAAAAAAAEAAAD///7//////////v/9//3//////////v/+//7//v/+//7//v////////8AAAAAAQACAAQABQAGAAcACAAKAAoACgALAAwADAAMAAwADAAMAAwADAAMAAwACwAKAAkACQAJAAgACAAIAAgABwAHAAYABQAEAAMAAgAAAP///v/9//z/+v/5//f/9f/0//P/8v/x//D/7//u/+3/7f/s/+z/7P/s/+7/7//w//L/8//0//X/9v/3//j/+v/8//3///8BAAMABgAIAAoADAAOABAAEQARABIAEwAUABQAFQAVABYAFAAUABYAFAAUABQAFgAVABMAFQAVABUAFAAVABUAFQAVABQAFAATABQAFAAUABQAFAATABQAFAAUABMAEgASABAAEAAQAA4ADgAMAAkACAAGAAIAAAD8//r/9//0//L/7//t/+v/6f/n/+b/4//h/+D/3v/e/9//3v/e/97/3v/e/+H/4v/j/+b/5v/n/+n/6v/s/+7/7//w//L/8//0//b/9//4//n/+v/7//z//f/+////AAABAAIABAAFAAcACAAJAAoACwANAA4ADwARABIAEwAUABUAFgAXABcAFwAXABcAFgAWABUAFAATABIAEAAPAA0ACwAJAAYABAABAP///f/7//n/9//1//T/8v/w/+//7f/s/+r/6f/o/+f/5v/l/+T/5P/j/+T/5P/l/+X/5v/n/+j/6P/p/+r/6//s/+z/7f/u/+7/7v/v/+//7//v/+//8P/w//H/8f/y//P/9P/1//f/+P/5//v//P/+/wEAAwAFAAcACgALAA4AEAATABUAFwAZABsAHQAfACEAIwAkACUAJgAmACUAJQAkACMAIQAgAB0AGwAaABkAGAAWABQAEgAPAA4ADAALAAgABgAEAAIAAQAAAAAA///+//7//f/8//v/+f/5//j/+P/4//f/9//2//X/9f/0//T/9P/0//X/9f/1//X/9f/2//f/9//4//j/+f/6//v//P/8//3//v///wAAAAABAAIAAwAEAAUABQAGAAcACAAIAAcACAAJAAgACAAIAAgABwAIAAkACQAIAAoACgAIAAkABwAKAAcABgAKAAgABwAKAAoACQAJAAoACwALAAwACgAKAAsACQAIAAgACAAGAAQAAwACAAEA///+//3//P/7//r/+f/4//j/9//3//f/9//3//f/9//3//j/+P/4//j/+P/5//n/+f/5//n/+f/5//j/+f/5//n/+f/4//n/+f/6//r/+f/6//v/+//6//r/+//8//3//P/8//7/AAABAAIAAgADAAQABQAGAAcACAAJALine truncated
+          }
+        }
+      ]
+    },
+    {
+      "role": "user",
+      "content": [
+        {
+          "type": "audio_url",
+          "audio_url": {
+            "url": "data:audio/wav;base64,UklGRkqYAABXQVZFZm10IBAAAAABAAEAwF0AAIC7AAACABAAZGF0YSaYAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAAAAAP//AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD//wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP//AAAAAP//AAAAAAAA///o/9T/xv/G/8b/yP/K/8r/y//M/8z/zf/O/83/zf/P/87/zv/O/87/zv/P/8//z//P/8//z//P/8//z//P/8//z//P/8//0P/Q/8//z//Q/9D/0f/R/9D/0f/R/9H/0f/R/9H/0f/R/9H/0f/S/9L/0v/S/9L/0v/S/9L/0v/T/9P/0//T/9P/0//T/9P/0//T/9P/0//T/9P/0//U/9T/1f/V/9T/1P/U/9X/1P/V/9X/1f/V/9b/1v/W/9b/1v/W/9b/1v/W/9f/1//X/9f/1//X/9f/1//X/9j/2P/Y/9j/2P/Y/9j/2P/Y/9j/2f/Z/9r/2v/a/9r/2v/b/9v/2//b/9v/2//c/9v/3P/c/9z/3P/d/9z/3P/c/93/3f/e/93/3f/e/97/3v/e/97/3//f/+D/4P/g/9//3//g/+D/4P/g/+H/4f/g/+H/4P/h/+H/4f/h/+L/4v/i/+H/4f/i/+L/4v/i/+L/4v/j/+L/4//j/+P/4//j/+P/4//j/+P/5P/k/+T/4//l/+T/5P/k/+T/5P/l/+T/5f/l/+X/5f/m/+X/5v/m/+b/5v/m/+b/5v/m/+b/5v/n/+b/5//n/+b/5v/o/+f/5//n/+j/5//p/+f/5//o/+j/6P/n/+j/6P/p/+n/6f/o/+j/6P/p/+j/6P/o/+n/6P/p/+j/6P/p/+n/6f/q/+n/6v/p/+n/6f/q/+n/6f/p/+n/6f/p/+n/6f/p/+n/6v/p/+n/6f/p/+n/6f/p/+n/6f/r/+r/6v/q/+r/6v/r/+n/6v/p/+r/6v/r/+n/6v/q/+r/6v/q/+r/6v/q/+r/6v/q/+r/6//q/+v/6//r/+r/6//r/+r/6v/q/+r/6v/q/+r/6v/q/+r/6//q/+v/6//s/+r/6//q/+r/6v/s/+n/6v/q/+v/6v/r/+v/6//s/+v/6//s/+v/7P/r/+r/6//s/+z/6//r/+v/7P/r/+z/6//r/+v/6//r/+r/6//r/+v/6//r/+v/6//s/+v/6//r/+z/7P/r/+v/6//s/+v/7P/r/+v/6//r/+v/6//s/+v/6//r/+z/7P/s/+z/7P/r/+z/6//s/+v/6//r/+v/6//s/+v/7P/r/+z/7P/s/+z/6//t/+v/7P/s/+v/7P/r/+z/6//s/+z/7P/r/+z/6//t/+v/7P/s/+v/7P/s/+z/6//s/+z/7P/s/+z/7P/s/+v/6//s/+z/7P/r/+z/7P/r/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/s/+v/6//s/+z/7P/s/+z/6//r/+z/7P/s/+v/7P/s/+v/7P/s/+v/6//s/+v/7P/s/+z/6//s/+v/7P/q/+v/6//s/+v/6//r/+z/6//r/+z/7P/r/+z/6//s/+z/7P/r/+z/7P/s/+v/6//t/+r/7P/r/+v/6//s/+v/7P/s/+v/7P/r/+z/6//s/+v/6//r/+z/6//r/+v/6//r/+v/6//t/+v/7P/r/+z/6//r/+z/7P/s/+v/6//r/+v/6//s/+v/7P/r/+z/7P/r/+v/6//s/+v/6//r/+z/6//s/+v/6//s/+v/7P/s/+z/7P/r/+v/6//s/+v/7P/r/+z/6//r/+v/6//s/+v/6//r/+z/7P/r/+z/6//s/+v/6//r/+v/6//s/+v/6//s/+r/6//r/+v/6//r/+v/6//r/+v/7P/r/+v/6//s/+v/6//r/+z/6//s/+v/6//q/+v/6//s/+v/6//s/+v/6//s/+v/6//r/+v/6//s/+v/6//r/+z/6//s/+v/6//s/+v/7P/q/+z/6//s/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/7P/r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//s/+r/6//r/+v/6//r/+r/6//r/+v/6//s/+v/6//q/+v/6v/s/+r/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+r/6//q/+v/6//r/+v/6//r/+v/6//r/+r/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/6//r/+z/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//s/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/6//s/+v/6//r/+z/6//r/+v/7P/r/+v/6//s/+v/6//q/+v/6//s/+r/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/6//s/+v/7P/r/+z/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+z/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//s/+v/6//s/+v/7P/r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6v/s/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/6//s/+z/7P/r/+v/7P/s/+v/6//s/+v/6//r/+v/6//r/+z/7P/s/+v/6//r/+z/6//s/+v/7P/r/+z/6//s/+v/6//s/+v/6//r/+z/6//s/+v/6//s/+v/7P/r/+v/6//s/+v/7P/r/+v/6//s/+v/7P/s/+v/7P/r/+z/7P/r/+z/6//s/+v/6//r/+z/6//s/+z/7P/s/+v/6//r/+v/6//r/+z/6//s/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+z/6//s/+v/6//s/+v/7P/s/+z/6//s/+v/6//s/+v/7P/s/+z/7P/s/+z/6//r/+v/6//s/+v/7P/s/+z/6//s/+v/6//r/+z/6//s/+v/6//r/+z/6//r/+v/6//s/+v/6//s/+v/6//r/+v/6//s/+v/7P/r/+z/7P/r/+v/7P/s/+z/7P/s/+v/7P/s/+z/6//s/+v/7P/s/+v/7P/s/+z/7P/s/+v/6//r/+v/6//s/+v/7P/r/+z/6//r/+z/6//s/+v/7P/r/+v/7P/r/+z/6//s/+v/6//s/+v/7P/r/+z/6//s/+z/7P/s/+v/7P/r/+z/7P/s/+v/6//r/+z/7P/s/+z/6//s/+v/6//s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/6//t/+v/7P/s/+z/7P/s/Line truncated
+          }
+        }
+      ]
+    }
+  ]
+}

examples/offline_inference/mimo_audio/message_convert.py

@@ -0,0 +1,714 @@
+import base64
+import io
+import os
+import random
+import re
+from collections.abc import Callable
+
+import librosa
+import numpy as np
+import torch
+import torchaudio
+from process_speechdata import InputSegment, StreamingInputSegment
+from torchaudio.transforms import MelSpectrogram
+
+speech_zeroemb_idx = 151667
+empty_token = "<|empty|>"
+mimo_audio_tokenizer = None
+device = "cpu"
+
+asr_zh_templates = [
+    "请将这段语音转换为文字",
+    "帮我识别这个音频文件中的内容",
+    "把这段录音转成文本",
+    "请转录这段语音",
+    "将音频内容转换成文字格式",
+    "识别并转写这段语音",
+    "把语音内容写成文字",
+    "转录这个音频片段",
+    "将这段对话转换为文本",
+    "麻烦帮我把这段录音整理成详细的文字记录",
+]
+
+asr_en_templates = [
+    "Please transcribe this audio file",
+    "Convert this speech recording to text",
+    "Transcribe the following voice message",
+    "Turn this audio into readable text",
+    "Please convert the recording to written format",
+    "Transcribe what you hear in this audio",
+    "Convert this spoken content to text",
+    "Please write down what is said in this recording",
+    "Transcribe this voice recording",
+    "Could you please help me transcribe this important recording?",
+    "Would you mind converting this voice message into a readable text format?",
+    "I'd really appreciate it if you could turn this audio file into a written document",
+]
+
+tts_zh_templates = [
+    "请将这段文字转换为语音",
+    "帮我把这个文本读出来",
+    "将这些文字生成音频",
+    "请朗读这段内容",
+    "把这段话转换成语音文件",
+    "生成这段文字的语音版本",
+    "请用语音播报这些内容",
+    "将文本转换为可听的音频",
+    "帮我朗读这段文字",
+    "把这些内容念出来",
+]
+
+tts_en_templates = [
+    "Please convert this text to speech",
+    "Turn this writing into audio",
+    "Generate speech from this text",
+    "Read this content out loud",
+    "Convert these words to voice",
+    "Create an audio version of this text",
+    "Please vocalize this content",
+    "Turn this text into audible format",
+    "Help me convert this writing to speech",
+    "Make this text into spoken audio",
+]
+
+
+def detect_language(text):
+    if re.search(r"[\u4e00-\u9fff]", text):
+        return "zh"
+    else:
+        return "en"
+
+
+# ============================================
+# Common Helper Functions - InputSegment Creation
+# ============================================
+
+
+def create_segment(text: str = "", audio=None) -> InputSegment:
+    """Create a standard InputSegment with default zeroemb parameters"""
+    return InputSegment(
+        text=text,
+        audio=audio,
+        speech_zeroemb_idx=speech_zeroemb_idx,
+        text_zeroemb_idx=empty_token,
+    )
+
+
+def create_streaming_segment(
+    text: str, audio, tokenizer, group_size: int, audio_channels: int
+) -> StreamingInputSegment:
+    """Create a StreamingInputSegment"""
+    return StreamingInputSegment(
+        text=text,
+        audio=audio,
+        tokenizer=tokenizer,
+        group_size=group_size,
+        audio_channels=audio_channels,
+        speech_zeroemb_idx=speech_zeroemb_idx,
+        text_zeroemb_idx=empty_token,
+    )
+
+
+# ============================================
+# Common Helper Functions - Common Token Segment Creation
+# ============================================
+
+
+def create_user_start() -> InputSegment:
+    """Create user role start token"""
+    return create_segment(text="<|im_start|>user\n")
+
+
+def create_user_end() -> InputSegment:
+    """Create role end token"""
+    return create_segment(text="<|im_end|>\n")
+
+
+def create_assistant_start() -> InputSegment:
+    """Create assistant role start token"""
+    return create_segment(text="<|im_start|>assistant\n")
+
+
+def create_system_start() -> InputSegment:
+    """Create system role start token"""
+    return create_segment(text="<|im_start|>system\n")
+
+
+def create_thinking_segment(thinking: bool = False) -> InputSegment:
+    """Create thinking token, closed or open based on thinking parameter"""
+    if thinking:
+        return create_segment(text="<think>\n")
+    else:
+        return create_segment(text="<think>\n\n</think>\n")
+
+
+def create_sostm_segment() -> InputSegment:
+    """Create streaming output start token"""
+    return create_segment(text="<|sostm|>")
+
+
+def create_assistant_start_with_sostm() -> InputSegment:
+    """Create assistant start token with sostm"""
+    return create_segment(text="<|im_start|>assistant\n<|sostm|>")
+
+
+def create_assistant_start_with_think() -> InputSegment:
+    """Create assistant start token with think opening"""
+    return create_segment(text="<|im_start|>assistant\n<think>\n")
+
+
+# ============================================
+# Common Helper Functions - Composite Segment Creation
+# ============================================
+
+
+def create_user_turn_with_audio(audio_tokenized, extra_text: str = None) -> list[InputSegment]:
+    """Create a user turn containing audio"""
+    segments = [
+        create_user_start(),
+        create_segment(audio=audio_tokenized),
+    ]
+    if extra_text:
+        segments.append(create_segment(text=extra_text))
+    segments.append(create_user_end())
+    return segments
+
+
+def create_user_turn_with_text(text: str) -> list[InputSegment]:
+    """Create a text-only user turn"""
+    return [
+        create_user_start(),
+        create_segment(text=text),
+        create_user_end(),
+    ]
+
+
+def create_system_turn_with_voice_prompt(prompt_text: str, audio_token) -> list[InputSegment]:
+    """Create a system turn with voice prompt"""
+    return [
+        create_system_start(),
+        create_segment(text=prompt_text),
+        create_segment(text="", audio=audio_token),
+        create_user_end(),
+    ]
+
+
+def create_system_turn_text_only(system_text: str) -> list[InputSegment]:
+    """Create a text-only system turn"""
+    return [
+        create_system_start(),
+        create_segment(text=system_text),
+        create_user_end(),
+    ]
+
+
+# ============================================
+# Common Helper Functions - Multi-turn Dialogue Processing
+# ============================================
+
+
+def process_multiturn_messages(
+    message_list: list[dict],
+    user_processor: Callable[[dict], list[InputSegment]],
+    assistant_processor: Callable[[dict], list[InputSegment]],
+) -> list[InputSegment]:
+    """
+    Generic multi-turn dialogue message processing function
+
+    Args:
+        message_list: List of messages, each containing 'role' and 'content'
+        user_processor: Function to process user messages
+        assistant_processor: Function to process assistant messages
+
+    Returns:
+        Processed list of InputSegments
+    """
+    lm_prompt = []
+    for message in message_list:
+        role = message["role"]
+        if role == "user":
+            lm_prompt.extend(user_processor(message))
+        elif role == "assistant":
+            lm_prompt.extend(assistant_processor(message))
+        else:
+            raise ValueError(f"Invalid role: {role}")
+    return lm_prompt
+
+
+def create_text_user_message(message: dict) -> list[InputSegment]:
+    """Process a text-only user message"""
+    return [
+        create_user_start(),
+        create_segment(text=message["content"]),
+        create_user_end(),
+    ]
+
+
+def create_text_assistant_message(message: dict) -> list[InputSegment]:
+    """Process a text-only assistant message"""
+    return [
+        create_assistant_start(),
+        create_segment(text=message["content"]),
+        create_user_end(),
+    ]
+
+
+def create_audio_user_message(message: dict) -> list[InputSegment]:
+    """Process an audio user message"""
+    return [
+        create_user_start(),
+        create_segment(audio=preprocess_input(message["content"])),
+        create_user_end(),
+    ]
+
+
+def append_assistant_ending(
+    lm_prompt: list[InputSegment], thinking: bool = False, use_sostm: bool = False
+) -> list[InputSegment]:
+    """
+    Append assistant ending to the prompt
+
+    Args:
+        lm_prompt: Existing prompt list
+        thinking: Whether to use open thinking token
+        use_sostm: Whether to use sostm token (for speech output)
+    """
+    if use_sostm:
+        lm_prompt.append(create_assistant_start_with_sostm())
+    else:
+        lm_prompt.append(create_assistant_start())
+        lm_prompt.append(create_thinking_segment(thinking))
+    return lm_prompt
+
+
+def get_asr_sft_prompt(
+    input: None | str = None,
+):
+    """Build prompt for ASR (Automatic Speech Recognition) task"""
+    audio_tokenized = preprocess_input(input)
+    template = random.choice(asr_zh_templates + asr_en_templates)
+
+    lm_prompt = create_user_turn_with_audio(audio_tokenized, extra_text=template)
+    lm_prompt = append_assistant_ending(lm_prompt, thinking=False)
+    return lm_prompt
+
+
+def resample_audio_if_needed(wav_tensor: torch.Tensor, original_sr: int):
+    target_sr = 24000
+    if original_sr != target_sr:
+        wav_tensor = torchaudio.functional.resample(wav_tensor, original_sr, target_sr)
+    return wav_tensor
+
+
+def wav2mel(wav, device="cpu"):
+    mel_transform = MelSpectrogram(
+        sample_rate=mimo_audio_tokenizer.config.sampling_rate,
+        n_fft=mimo_audio_tokenizer.config.nfft,
+        hop_length=mimo_audio_tokenizer.config.hop_length,
+        win_length=mimo_audio_tokenizer.config.window_size,
+        f_min=mimo_audio_tokenizer.config.fmin,
+        f_max=mimo_audio_tokenizer.config.fmax,
+        n_mels=mimo_audio_tokenizer.config.n_mels,
+        power=1.0,
+        center=True,
+    ).to(device)
+    spec = mel_transform(wav[None, :])
+    return torch.log(torch.clip(spec, min=1e-7)).squeeze()
+
+
+def group_by_length(features: torch.Tensor, lengths: torch.Tensor, max_length: int):
+    if features.size(0) != lengths.sum().item():
+        raise ValueError(f"Feature size mismatch: {features.size(0)} vs {lengths.sum().item()}")
+
+    split_points = []
+    current_sum = 0
+
+    for i, seq_len in enumerate(lengths):
+        if current_sum + seq_len > max_length and current_sum > 0:
+            split_points.append(i)
+            current_sum = seq_len.item()
+        else:
+            current_sum += seq_len.item()
+
+    # Convert split points to group sizes
+    group_sizes = []
+    prev = 0
+    for point in split_points:
+        group_sizes.append(point - prev)
+        prev = point
+    if prev < len(lengths):
+        group_sizes.append(len(lengths) - prev)
+
+    len_groups = torch.split(lengths, group_sizes)
+    feature_sizes = [group.sum().item() for group in len_groups]
+    feature_groups = torch.split(features, feature_sizes)
+
+    return feature_groups, len_groups
+
+
+def encode_batch(input_features: torch.Tensor, input_lens: torch.Tensor, max_length: int = 256000):
+    feature_groups, len_groups = group_by_length(input_features, input_lens, max_length)
+
+    encoded_parts = []
+    for features, lengths in zip(feature_groups, len_groups):
+        with torch.no_grad():
+            codes, _ = mimo_audio_tokenizer.encoder.encode(
+                input_features=features.to(device), input_lens=lengths.to(device), return_codes_only=True
+            )
+            encoded_parts.append(codes)
+
+    return torch.cat(encoded_parts, dim=-1)
+
+
+def preprocess_input(input: None | str | torch.Tensor = None, device="cpu", audio_channels=4, group_size=8):
+    if isinstance(input, torch.Tensor) or (isinstance(input, str) and os.path.isfile(input)):
+        return "<|sosp|><|empty|><|eosp|>"
+
+    else:
+        text = input
+        if (
+            text.isupper() or text.islower()
+        ):  # If the text only contains upper-case or lower-case letters, capitalize it.
+            text = text.capitalize()
+        return text
+
+
+def _build_tts_system_prompt(has_voice_prompt: bool, voice_audio_token=None) -> list[InputSegment]:
+    """Build system prompt for TTS task"""
+    if has_voice_prompt and voice_audio_token is not None:
+        return [
+            create_system_start(),
+            create_segment(
+                # text="You need to generate a speech with the same timbre as the speech prompt, based on the specified style instructions and text content. Your timbre should be: "
+                text="你需要根据指定的风格指令和文本内容来生成和语音prompt具有相同音色的语音。你的音色应该是："
+            ),
+            create_segment(text="", audio=voice_audio_token),
+            create_user_end(),
+        ]
+    else:
+        return create_system_turn_text_only("你需要根据指定的风格指令和文本内容来生成语音。")
+        # return create_system_turn_text_only(
+        #     "You need to generate speech based on the specified style instructions and text content."
+        # )
+
+
+def _build_tts_system_prompt_no_instruct(has_voice_prompt: bool, voice_audio_token=None) -> list[InputSegment]:
+    """Build system prompt for TTS task"""
+    if has_voice_prompt and voice_audio_token is not None:
+        return [
+            create_system_start(),
+            create_segment(
+                # text="You need to generate a speech with the same timbre as the speech prompt, based on the specified style instructions and text content. Your timbre should be:"
+                text="你需要根据指定的风格指令和文本内容来生成和语音prompt具有相同音色的语音。你的音色应该是："
+            ),
+            create_segment(text="", audio=voice_audio_token),
+            create_user_end(),
+        ]
+    else:
+        return []
+
+
+def get_tts_sft_prompt(
+    input: None | str = None,
+    instruct=None,
+    read_text_only=True,
+    prompt_speech=None,
+):
+    """
+    Build prompt for TTS (Text-to-Speech) task
+
+    Args:
+        input: Input text
+        instruct: Style instruction (e.g., "speak happily in a child's voice")
+        read_text_only: Whether to read only plain text (False means text contains template)
+        prompt_speech: Reference audio (for voice cloning)
+    """
+    assistant_prompt_audio_token = preprocess_input(prompt_speech) if prompt_speech is not None else None
+
+    if not read_text_only:
+        # Not just reading text, text contains template (template:text format)
+        text = preprocess_input(input)
+        lm_prompt = _build_tts_system_prompt(
+            has_voice_prompt=assistant_prompt_audio_token is not None,
+            voice_audio_token=assistant_prompt_audio_token,
+        )
+        lm_prompt.append(create_segment(text=f"<|im_start|>user\n{text}<|im_end|>\n"))
+        lm_prompt.append(create_assistant_start_with_think())
+    else:
+        # Plain text (no instruction inside)
+        language = detect_language(input)
+        template = random.choice(tts_zh_templates if language == "zh" else tts_en_templates)
+        text = preprocess_input(input)
+
+        if instruct is None:
+            # No instruct instruction
+            lm_prompt = _build_tts_system_prompt_no_instruct(
+                has_voice_prompt=assistant_prompt_audio_token is not None,
+                voice_audio_token=assistant_prompt_audio_token,
+            )
+            lm_prompt.extend(
+                [
+                    create_segment(text=f"<|im_start|>user\n{template}: {text}<|im_end|>\n"),
+                    create_assistant_start_with_sostm(),
+                ]
+            )
+        else:
+            # Has instruct instruction
+            lm_prompt = _build_tts_system_prompt(
+                has_voice_prompt=assistant_prompt_audio_token is not None,
+                voice_audio_token=assistant_prompt_audio_token,
+            )
+            lm_prompt.append(create_segment(text=f"<|im_start|>user\n{template}: {text}({instruct})<|im_end|>\n"))
+            lm_prompt.append(create_assistant_start_with_think())
+
+    return lm_prompt
+
+
+def get_audio_understanding_sft_prompt(
+    input_speech,
+    input_text,
+    thinking=False,
+    use_sostm=False,
+):
+    """Build prompt for audio understanding task"""
+    audio_tokenized = preprocess_input(input_speech)
+
+    lm_prompt = create_user_turn_with_audio(audio_tokenized, extra_text=input_text)
+    lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking, use_sostm=use_sostm)
+    return lm_prompt
+
+
+def _build_voice_prompt_system(prompt_speech) -> list[InputSegment]:
+    """Build system prompt with voice prompt"""
+    return create_system_turn_with_voice_prompt(
+        prompt_text="Your voice should be：", audio_token=preprocess_input(prompt_speech)
+    )
+
+
+def get_spoken_dialogue_sft_prompt(
+    input_speech,
+    system_prompt=None,
+    prompt_speech=None,
+    add_history=False,
+):
+    """
+    Build prompt for spoken dialogue task
+
+    Args:
+        input_speech: Input speech
+        system_prompt: System prompt text
+        prompt_speech: Reference audio (for voice cloning)
+        add_history: Whether to add history (Note: history variable is undefined in original code)
+    """
+    audio_tokenized = preprocess_input(input_speech)
+    lm_prompt = []
+
+    # Note: history variable is undefined in original code, this branch may never execute
+    # To use history feature, history should be passed as a parameter
+    if add_history:
+        # Simplified form of adding history
+        lm_prompt = create_user_turn_with_audio(audio_tokenized)
+        lm_prompt.append(create_assistant_start_with_sostm())
+    else:
+        # Add voice prompt (if available)
+        if prompt_speech:
+            lm_prompt.extend(_build_voice_prompt_system(prompt_speech))
+
+        # Add user turn
+        lm_prompt.append(create_user_start())
+        if system_prompt:
+            lm_prompt.append(create_segment(text=system_prompt))
+        lm_prompt.append(create_segment(audio=audio_tokenized))
+        lm_prompt.append(create_user_end())
+        lm_prompt.append(create_assistant_start_with_sostm())
+
+    return lm_prompt
+
+
+def get_spoken_dialogue_sft_multiturn_prompt(
+    message_list,
+    system_prompt=None,
+    prompt_speech=None,
+    tokenizer=None,
+    group_size=8,
+    audio_channels=4,
+):
+    """
+    Build prompt for multi-turn spoken dialogue task
+
+    Args:
+        message_list: List of messages containing role and content
+        system_prompt: System prompt text
+        prompt_speech: Reference audio (for voice cloning)
+        tokenizer: Tokenizer
+        group_size: Group size
+        audio_channels: Number of audio channels
+    """
+    lm_prompt = []
+
+    # Add voice prompt (if available)
+    if prompt_speech:
+        lm_prompt.extend(
+            create_system_turn_with_voice_prompt(
+                prompt_text="Your voice should be:", audio_token=preprocess_input(prompt_speech)
+            )
+        )
+
+    # Define message processors
+    def user_processor(msg):
+        segments = [create_user_start()]
+        if system_prompt:
+            segments.append(create_segment(text=system_prompt))
+        segments.append(create_segment(audio=preprocess_input(msg["content"])))
+        segments.append(create_user_end())
+        return segments
+
+    def assistant_processor(msg):
+        return [
+            create_assistant_start(),
+            create_streaming_segment(
+                text=msg["content"]["text"],
+                audio=preprocess_input(msg["content"]["audio"]),
+                tokenizer=tokenizer,
+                group_size=group_size,
+                audio_channels=audio_channels,
+            ),
+            create_user_end(),
+        ]
+
+    # Process message list
+    lm_prompt.extend(process_multiturn_messages(message_list, user_processor, assistant_processor))
+    lm_prompt.append(create_assistant_start_with_sostm())
+
+    return lm_prompt
+
+
+def get_s2t_dialogue_sft_prompt(
+    input_speech,
+    thinking=False,
+):
+    """Build prompt for speech-to-text dialogue task"""
+    audio_tokenized = preprocess_input(input_speech)
+
+    lm_prompt = create_user_turn_with_audio(audio_tokenized)
+    lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
+    return lm_prompt
+
+
+def get_s2t_dialogue_sft_multiturn_prompt(message_list, thinking=False):
+    """Build prompt for multi-turn speech-to-text dialogue task"""
+    lm_prompt = process_multiturn_messages(
+        message_list, user_processor=create_audio_user_message, assistant_processor=create_text_assistant_message
+    )
+    lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
+    return lm_prompt
+
+
+def get_text_dialogue_sft_prompt(
+    input_text,
+    thinking=False,
+):
+    """Build prompt for text-only dialogue task"""
+    lm_prompt = create_user_turn_with_text(input_text)
+    lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
+    return lm_prompt
+
+
+def get_text_dialogue_sft_multiturn_prompt(
+    message_list,
+    thinking=False,
+):
+    """Build prompt for multi-turn text-only dialogue task"""
+    lm_prompt = process_multiturn_messages(
+        message_list, user_processor=create_text_user_message, assistant_processor=create_text_assistant_message
+    )
+    lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
+    return lm_prompt
+
+
+def get_in_context_learning_s2s_prompt(
+    instruction,
+    prompt_examples,
+    audio,
+    tokenizer=None,
+    group_size=8,
+    audio_channels=4,
+):
+    """
+    Build prompt for In-Context Learning speech-to-speech task
+
+    Args:
+        instruction: Instruction text
+        prompt_examples: List of examples, each containing input_audio, output_transcription, output_audio
+        audio: Input audio to be processed
+        tokenizer: Tokenizer
+        group_size: Group size
+        audio_channels: Number of audio channels
+    """
+    prompt = [create_segment(text=f"[Int]:{instruction}\n")]
+
+    # Add examples
+    for example in prompt_examples:
+        prompt.extend(
+            [
+                create_segment(audio=preprocess_input(example["input_audio"])),
+                create_segment(text="\n"),
+                create_streaming_segment(
+                    text=example["output_transcription"],
+                    audio=preprocess_input(example["output_audio"]),
+                    tokenizer=tokenizer,
+                    group_size=group_size,
+                    audio_channels=audio_channels,
+                ),
+                create_segment(text=" \n\n"),
+            ]
+        )
+
+    # Add input audio to be processed
+    prompt.extend(
+        [
+            create_segment(audio=preprocess_input(audio)),
+            create_segment(text="\n"),
+            create_sostm_segment(),
+        ]
+    )
+
+    return prompt
+
+
+def get_audio_data(audio_url):
+    if audio_url.startswith("data:"):
+        header, b64_data = audio_url.split(",", 1)
+        audio_bytes = base64.b64decode(b64_data.strip())
+        audio_file = io.BytesIO(audio_bytes)
+    else:
+        # File path
+        audio_file = audio_url
+
+    audio_signal, sr = librosa.load(audio_file, sr=24000)
+    audio_data = (audio_signal.astype(np.float32), sr)
+    return audio_data
+
+
+def to_prompt(input_segs):
+    out_put = []
+
+    for input_seg in input_segs:
+        if isinstance(input_seg, StreamingInputSegment) and input_seg.text:
+            out_put.append("<|sostm|>")
+            if input_seg.audio is not None and isinstance(input_seg.audio, str):
+                out_put.append(input_seg.text)
+                out_put.append("<|eot|>")
+                out_put.append("<|empty|>")
+            else:
+                out_put.append(input_seg.text)
+                out_put.append("<|eot|>")
+            out_put.append("<|eostm|>")
+
+        else:
+            out_put.append(input_seg.text)
+            if input_seg.audio is not None:
+                out_put.append(input_seg.audio)
+
+    prompt = "".join(out_put)
+    return prompt

examples/offline_inference/mimo_audio/process_speechdata.py

@@ -0,0 +1,261 @@
+#!/usr/bin/env python3
+# Copyright 2025 Xiaomi Corporation.
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+
+
+import torch
+
+
+class InputSegment:
+    def __init__(
+        self,
+        text: str = "",
+        audio: torch.Tensor = None,
+        tokenized_text: torch.Tensor = None,
+        speech_zeroemb_idx: int | list[int] = 1024,
+        text_zeroemb_idx: int = 152067,
+        add_sosp_eosp=True,
+    ) -> None:
+        has_text = text is not None
+        has_tokenized_text = tokenized_text is not None
+        assert has_text or has_tokenized_text, "Text or tokenized text must be provided"
+
+        self.audio = audio
+        self.text = text
+        self.tokenized_text = tokenized_text
+        self.speech_zeroemb_idx = speech_zeroemb_idx
+        self.text_zeroemb_idx = text_zeroemb_idx
+        self.add_sosp_eosp = add_sosp_eosp
+
+    @staticmethod
+    def insert_between(tensor, i, value=-1):
+        return torch.scatter(
+            torch.full(
+                (1, tensor.shape[1] + (tensor.shape[1] - 1) * i + i),
+                value,
+                dtype=tensor.dtype,
+            ),
+            1,
+            torch.arange(0, tensor.shape[1], dtype=torch.int64)[None] * (i + 1),
+            tensor,
+        )
+
+    def to_input_id(
+        self,
+        tokenizer,
+        group_size: int,
+        audio_channels: int = 8,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if (
+            self.audio is None
+        ):  ## If no audio is provided, tokenize the text with (group_size-1) tokens inserted between each, fill with -100; and set audio to empty indices
+            if self.tokenized_text is None:
+                tokenized_text = tokenizer(
+                    self.text,
+                    return_tensors="pt",
+                    truncation=True,
+                    max_length=999999,
+                    padding=False,
+                    add_special_tokens=False,
+                )["input_ids"].int()
+            else:
+                tokenized_text = self.tokenized_text.unsqueeze(0)
+
+            if group_size > 1:
+                tokenized_text = self.insert_between(tokenized_text, group_size - 1, value=-100)
+
+            if isinstance(self.speech_zeroemb_idx, list):
+                audio_part_input_id = torch.zeros((audio_channels, tokenized_text.shape[1]), dtype=torch.int)
+                for i, idx in enumerate(self.speech_zeroemb_idx):
+                    audio_part_input_id[i, :] = idx
+            else:
+                audio_part_input_id = torch.full(
+                    (audio_channels, tokenized_text.shape[1]), self.speech_zeroemb_idx, dtype=torch.int
+                )
+
+        else:  # If audio is provided, add sosp/eosp markers and convert audio: fill with empty tokens based on audio length
+            sosp_token = tokenizer.convert_tokens_to_ids("<|sosp|>") if self.add_sosp_eosp else None
+            eosp_token = tokenizer.convert_tokens_to_ids("<|eosp|>") if self.add_sosp_eosp else None
+            audio_part = self.audio.reshape(-1, audio_channels).T  # [audio_channels, seqlen]
+
+            assert audio_part.shape[1] % group_size == 0, (
+                f"Audio shape {audio_part.shape} is not divisible by group_size {group_size}"
+            )
+
+            text_len = audio_part.shape[1] // group_size
+            empty_token = self.text_zeroemb_idx
+            if empty_token is None:
+                empty_token = tokenizer.eod
+            tokenized_text = torch.full((1, text_len), empty_token, dtype=torch.int)
+
+            tokenized_text = (
+                torch.cat(
+                    [
+                        torch.tensor([[sosp_token]], dtype=torch.int),
+                        tokenized_text,
+                        torch.tensor([[eosp_token]], dtype=torch.int),
+                    ],
+                    dim=1,
+                )
+                if self.add_sosp_eosp
+                else tokenized_text
+            )
+            tokenized_text = self.insert_between(tokenized_text, group_size - 1, value=-100)
+
+            if self.add_sosp_eosp:
+                if isinstance(self.speech_zeroemb_idx, list):
+                    sosp_part = torch.zeros((audio_channels, group_size), dtype=torch.int)
+                    eosp_part = torch.zeros((audio_channels, group_size), dtype=torch.int)
+                    for i, idx in enumerate(self.speech_zeroemb_idx):
+                        sosp_part[i, :] = idx
+                        eosp_part[i, :] = idx
+                    audio_part_input_id = torch.cat([sosp_part, audio_part, eosp_part], dim=1)
+                else:
+                    audio_part_input_id = torch.cat(
+                        [
+                            torch.full((audio_channels, group_size), self.speech_zeroemb_idx, dtype=torch.int),
+                            audio_part,
+                            torch.full((audio_channels, group_size), self.speech_zeroemb_idx, dtype=torch.int),
+                        ],
+                        dim=1,
+                    )
+            else:
+                audio_part_input_id = audio_part
+
+        input_ids = torch.cat([tokenized_text, audio_part_input_id], dim=0)  # [n_rvq + 1, seqlen]
+
+        return input_ids
+
+
+class StreamingInputSegment:
+    def __init__(
+        self,
+        text: str = "",
+        audio: torch.Tensor = None,
+        tokenized_text: torch.Tensor = None,
+        speech_zeroemb_idx: int | list[int] = 1024,
+        text_zeroemb_idx: int = 152067,
+        text_segment_size: int = 5,
+        audio_segment_size: int = 5,
+        tokenizer=None,
+        group_size=None,
+        audio_channels=None,
+    ) -> None:
+        has_text = text is not None
+        has_tokenized_text = tokenized_text is not None
+        assert has_text or has_tokenized_text, "Text or tokenized text must be provided"
+
+        self.audio = audio
+        self.text = text
+        self.tokenized_text = tokenized_text
+        self.speech_zeroemb_idx = speech_zeroemb_idx
+        self.text_zeroemb_idx = text_zeroemb_idx
+        self.text_segment_size = text_segment_size
+        self.audio_segment_size = audio_segment_size
+        self.tokenizer = tokenizer
+        self.group_size = group_size
+        self.audio_channels = audio_channels
+
+    def to_input_id(
+        self,
+        tokenizer,
+        group_size: int,
+        audio_channels: int = 8,
+    ):
+        if self.tokenized_text is None:
+            tokenized_text = tokenizer(
+                self.text,
+                return_tensors="pt",
+                truncation=True,
+                max_length=999999,
+                padding=False,
+                add_special_tokens=False,
+            )["input_ids"].int()  # [1, seqlen]
+        else:
+            tokenized_text = self.tokenized_text.unsqueeze(0)
+
+        tokenized_text = tokenized_text.squeeze(0)
+
+        text_segments = tokenized_text.split(self.text_segment_size, dim=0)
+        audio_segments = self.audio.split(self.audio_segment_size * group_size * audio_channels, dim=0)
+
+        tokenized_segments = []
+        tokenized_segments.append(
+            InputSegment(
+                text="<|sostm|>",
+                speech_zeroemb_idx=self.speech_zeroemb_idx,
+                text_zeroemb_idx=self.text_zeroemb_idx,
+            ),
+        )
+
+        eot_tokens = tokenizer(
+            "<|eot|>",
+            return_tensors="pt",
+            truncation=True,
+            max_length=999999,
+            padding=False,
+            add_special_tokens=False,
+        )["input_ids"][0].to(text_segments[-1])
+
+        text_segments = text_segments[:-1] + (torch.cat([text_segments[-1], eot_tokens], dim=0),)
+
+        length = min(len(text_segments), len(audio_segments))  # The number of empty tokens here may vary
+        for i in range(length):
+            text_segment = text_segments[i]
+            audio_segment = audio_segments[i]
+
+            tokenized_segments.append(
+                InputSegment(
+                    tokenized_text=text_segment,
+                    speech_zeroemb_idx=self.speech_zeroemb_idx,
+                    text_zeroemb_idx=self.text_zeroemb_idx,
+                ),
+            )
+            tokenized_segments.append(
+                InputSegment(
+                    audio=audio_segment,
+                    add_sosp_eosp=False,
+                    speech_zeroemb_idx=self.speech_zeroemb_idx,
+                    text_zeroemb_idx=self.text_zeroemb_idx,
+                ),
+            )
+
+        for j in range(length, len(text_segments)):
+            tokenized_segments.append(
+                InputSegment(
+                    tokenized_text=text_segments[j],
+                    speech_zeroemb_idx=self.speech_zeroemb_idx,
+                    text_zeroemb_idx=self.text_zeroemb_idx,
+                ),
+            )
+
+        for j in range(length, len(audio_segments)):
+            tokenized_segments.append(
+                InputSegment(
+                    audio=audio_segments[j],
+                    add_sosp_eosp=False,
+                    speech_zeroemb_idx=self.speech_zeroemb_idx,
+                    text_zeroemb_idx=self.text_zeroemb_idx,
+                ),
+            )
+
+        tokenized_segments.append(
+            InputSegment(
+                text="<|eostm|>",
+                speech_zeroemb_idx=self.speech_zeroemb_idx,
+                text_zeroemb_idx=self.text_zeroemb_idx,
+            ),
+        )
+
+        input_ids = [
+            seg.to_input_id(
+                self.tokenizer,
+                self.group_size,
+                self.audio_channels,
+            )
+            for seg in tokenized_segments
+        ]
+
+        input_ids = torch.cat(input_ids, dim=1).type(torch.int64)  # [n_rvq + 1, seqlen]
+
+        return input_ids

examples/offline_inference/qwen3_tts/README.md

@@ -87,6 +87,19 @@ Examples:
 python end2end.py --query-type Base --mode-tag icl
 ```
 
+## Streaming Mode
+
+Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config):
+
+```bash
+python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_stream
+```
+
+Each 25-frame Code2Wav chunk is logged as it arrives. The final WAV file is written once generation
+completes. This demonstrates that audio data is available progressively rather than only at the end.
+
+> **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged.
+
 ## Batched Decoding
 
 The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_batch_size > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.

examples/offline_inference/qwen3_tts/end2end.py

@@ -4,6 +4,7 @@ Provides single and batch sample inputs for CustomVoice, VoiceDesign, and Base
 tasks, then runs Omni generation and saves output wav files.
 """
 
+import asyncio
 import logging
 import os
 from typing import Any, NamedTuple
@@ -15,7 +16,7 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
-from vllm_omni import Omni
+from vllm_omni import AsyncOmni, Omni
 
 logger = logging.getLogger(__name__)
 
@@ -242,12 +243,15 @@ def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> Que
     )
 
 
-def main(args):
-    """Run offline inference with Omni using prepared sample inputs.
+query_map = {
+    "CustomVoice": get_custom_voice_query,
+    "VoiceDesign": get_voice_design_query,
+    "Base": get_base_query,
+}
 
-    Args:
-        args: Parsed CLI args from parse_args().
-    """
+
+def _build_inputs(args) -> tuple[str, list]:
+    """Resolve model name and inputs list from CLI args."""
     if args.batch_size < 1 or (args.batch_size & (args.batch_size - 1)) != 0:
         raise ValueError(
             f"--batch-size must be a power of two (got {args.batch_size}); "
@@ -259,40 +263,49 @@ def main(args):
     if args.query_type in {"CustomVoice", "VoiceDesign"}:
         query_result = query_func(use_batch_sample=args.use_batch_sample)
     elif args.query_type == "Base":
-        query_result = query_func(
-            use_batch_sample=args.use_batch_sample,
-            mode_tag=args.mode_tag,
-        )
+        query_result = query_func(use_batch_sample=args.use_batch_sample, mode_tag=args.mode_tag)
     else:
         query_result = query_func()
 
     model_name = query_result.model_name
 
-    # Load prompts from text file if provided.
-    # Use the default query as a template so task-specific fields
-    # (e.g. ref_audio for Base) are preserved; only override text.
     if args.txt_prompts:
         with open(args.txt_prompts) as f:
             lines = [line.strip() for line in f if line.strip()]
         if not lines:
             raise ValueError(f"No valid prompts found in {args.txt_prompts}")
-        template = query_result.inputs
-        if isinstance(template, list):
-            template = template[0]
+        template = query_result.inputs if not isinstance(query_result.inputs, list) else query_result.inputs[0]
         template_info = template["additional_information"]
-        inputs = []
-        for text in lines:
-            additional_information = {**template_info, "text": [text]}
-            inputs.append(
-                {
-                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
-                    "additional_information": additional_information,
-                }
-            )
+        inputs = [
+            {
+                "prompt_token_ids": [0] * _estimate_prompt_len({**template_info, "text": [t]}, model_name),
+                "additional_information": {**template_info, "text": [t]},
+            }
+            for t in lines
+        ]
     else:
-        inputs = query_result.inputs
-        if not isinstance(inputs, list):
-            inputs = [inputs]
+        inputs = query_result.inputs if isinstance(query_result.inputs, list) else [query_result.inputs]
+
+    return model_name, inputs
+
+
+def _save_wav(output_dir: str, request_id: str, mm: dict) -> None:
+    """Concatenate audio chunks and write to a wav file."""
+    audio_data = mm["audio"]
+    sr_raw = mm["sr"]
+    sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
+    sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
+    audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
+    out_wav = os.path.join(output_dir, f"output_{request_id}.wav")
+    sf.write(out_wav, audio_tensor.float().cpu().numpy().flatten(), samplerate=sr, format="WAV")
+    logger.info(f"Request ID: {request_id}, Saved audio to {out_wav}")
+
+
+def main(args):
+    """Run offline inference with Omni."""
+    model_name, inputs = _build_inputs(args)
+    output_dir = args.output_dir
+    os.makedirs(output_dir, exist_ok=True)
 
     omni = Omni(
         model=model_name,
@@ -301,43 +314,40 @@ def main(args):
         stage_init_timeout=args.stage_init_timeout,
     )
 
-    output_dir = args.output_dir
-    os.makedirs(output_dir, exist_ok=True)
-
     batch_size = args.batch_size
     for batch_start in range(0, len(inputs), batch_size):
         batch = inputs[batch_start : batch_start + batch_size]
-        omni_generator = omni.generate(batch, sampling_params_list=None)
-        for stage_outputs in omni_generator:
+        for stage_outputs in omni.generate(batch):
             for output in stage_outputs.request_output:
-                request_id = output.request_id
-                audio_data = output.outputs[0].multimodal_output["audio"]
-                # async_chunk mode returns a list of chunks; concatenate them.
-                if isinstance(audio_data, list):
-                    audio_tensor = torch.cat(audio_data, dim=-1)
-                else:
-                    audio_tensor = audio_data
-                output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
-                sr_val = output.outputs[0].multimodal_output["sr"]
-                audio_samplerate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val[-1])
-                # Convert to numpy array and ensure correct format
-                audio_numpy = audio_tensor.float().detach().cpu().numpy()
+                _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
 
-                # Ensure audio is 1D (flatten if needed)
-                if audio_numpy.ndim > 1:
-                    audio_numpy = audio_numpy.flatten()
 
-                # Save audio file with explicit WAV format
-                sf.write(output_wav, audio_numpy, samplerate=audio_samplerate, format="WAV")
-                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+async def main_streaming(args):
+    """Run offline inference with AsyncOmni, logging each audio chunk as it arrives."""
+    model_name, inputs = _build_inputs(args)
+    output_dir = args.output_dir
+    os.makedirs(output_dir, exist_ok=True)
+
+    omni = AsyncOmni(
+        model=model_name,
+        stage_configs_path=args.stage_configs_path,
+        log_stats=args.log_stats,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+
+    for i, prompt in enumerate(inputs):
+        request_id = str(i)
+        async for stage_output in omni.generate(prompt, request_id=request_id):
+            mm = stage_output.request_output.outputs[0].multimodal_output
+            if not stage_output.finished:
+                audio = mm.get("audio")
+                n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1)
+                logger.info(f"Request {request_id}: received chunk {n}")
+            else:
+                _save_wav(output_dir, request_id, mm)
 
 
 def parse_args():
-    """Parse CLI arguments for offline TTS inference.
-
-    Returns:
-        argparse.Namespace with CLI options.
-    """
     parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
     parser.add_argument(
         "--query-type",
@@ -438,6 +448,12 @@ def parse_args():
         choices=["icl", "xvec_only"],
         help="Mode tag for Base query x_vector_only_mode (default: icl).",
     )
+    parser.add_argument(
+        "--streaming",
+        action="store_true",
+        default=False,
+        help="Stream audio chunks as they arrive via AsyncOmni (async_chunk mode only).",
+    )
     parser.add_argument(
         "--batch-size",
         type=int,
@@ -448,13 +464,9 @@ def parse_args():
     return parser.parse_args()
 
 
-query_map = {
-    "CustomVoice": get_custom_voice_query,
-    "VoiceDesign": get_voice_design_query,
-    "Base": get_base_query,
-}
-
-
 if __name__ == "__main__":
     args = parse_args()
-    main(args)
+    if args.streaming:
+        asyncio.run(main_streaming(args))
+    else:
+        main(args)

examples/offline_inference/text_to_image/text_to_image.py

@@ -132,10 +132,18 @@ def parse_args() -> argparse.Namespace:
         "--quantization",
         type=str,
         default=None,
-        choices=["fp8"],
-        help="Quantization method for the transformer. "
-        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs). "
-        "Default: None (no quantization, uses BF16).",
+        choices=["fp8", "gguf"],
+        help=(
+            "Quantization method for the transformer. "
+            "Options: 'fp8' (FP8 W8A8), 'gguf' (GGUF quantized weights). "
+            "Default: None (no quantization, uses BF16)."
+        ),
+    )
+    parser.add_argument(
+        "--gguf-model",
+        type=str,
+        default=None,
+        help=("GGUF file path or HF reference for transformer weights. Required when --quantization gguf is set."),
     )
     parser.add_argument(
         "--ignored-layers",
@@ -174,16 +182,6 @@ def parse_args() -> argparse.Namespace:
         default=1.0,
         help="Scale factor for LoRA weights (default: 1.0).",
     )
-    parser.add_argument(
-        "--vae_use_slicing",
-        action="store_true",
-        help="Enable VAE slicing for memory optimization.",
-    )
-    parser.add_argument(
-        "--vae_use_tiling",
-        action="store_true",
-        help="Enable VAE tiling for memory optimization.",
-    )
     parser.add_argument(
         "--vae-patch-parallel-size",
         type=int,
@@ -275,7 +273,14 @@ def main():
     # ignored_layers is specified so the list flows through OmniDiffusionConfig
     quant_kwargs: dict[str, Any] = {}
     ignored_layers = [s.strip() for s in args.ignored_layers.split(",") if s.strip()] if args.ignored_layers else None
-    if args.quantization and ignored_layers:
+    if args.quantization == "gguf":
+        if not args.gguf_model:
+            raise ValueError("--gguf-model is required when --quantization gguf is set.")
+        quant_kwargs["quantization_config"] = {
+            "method": "gguf",
+            "gguf_model": args.gguf_model,
+        }
+    elif args.quantization and ignored_layers:
         quant_kwargs["quantization_config"] = {
             "method": args.quantization,
             "ignored_layers": ignored_layers,
@@ -286,7 +291,6 @@ def main():
     omni_kwargs = {
         "model": args.model,
         "enable_layerwise_offload": args.enable_layerwise_offload,
-        "layerwise_num_gpu_layers": args.layerwise_num_gpu_layers,
         "vae_use_slicing": args.vae_use_slicing,
         "vae_use_tiling": args.vae_use_tiling,
         "cache_backend": args.cache_backend,

examples/online_serving/glm_image/README.md

@@ -0,0 +1,247 @@
+# GLM-Image Online Serving
+
+This example demonstrates how to deploy GLM-Image for online image generation using vLLM-Omni.
+
+## 🛠️ Installation
+
+Please refer to [README.md](../../../README.md)
+
+## Run examples (GLM-Image)
+
+**Note**: These examples work with the default configuration on **2× NVIDIA A100 (80GB)** or equivalent. Stage 0 (AR) and Stage 1 (Diffusion) each use one GPU by default. For single-GPU setups, modify the stage configuration to share the same device.
+
+### Launch the Server
+
+```bash
+# Use default configuration
+vllm serve zai-org/GLM-Image --omni --port 8091
+```
+
+Or use the convenience script:
+
+```bash
+cd examples/online_serving/glm_image
+bash run_server.sh
+```
+
+If you have a custom stage configs file:
+
+```bash
+vllm serve zai-org/GLM-Image --omni --port 8091 --stage-configs-path /path/to/glm_image.yaml
+```
+
+### Send Requests
+
+Get into the glm_image folder:
+
+```bash
+cd examples/online_serving/glm_image
+```
+
+Send request via Python:
+
+```bash
+python openai_chat_client.py --prompt "A cute cat sitting on a window sill"
+```
+
+The Python client supports the following command-line arguments:
+
+- `--prompt` (or `-p`): Text prompt for generation (default: `A beautiful sunset over the ocean with sailing boats`)
+- `--output` (or `-o`): Output file path (default: `glm_image_output.png`)
+- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
+- `--image` (or `-i`): Input image path (for image-to-image editing)
+- `--height`: Image height in pixels (default: 1024)
+- `--width`: Image width in pixels (default: 1024)
+- `--steps`: Number of inference steps (default: 50)
+- `--guidance-scale`: Classifier-free guidance scale (default: 1.5)
+- `--seed`: Random seed (default: 42)
+- `--negative`: Negative prompt
+
+## Modality Control
+
+GLM-Image supports **text-to-image** and **image-to-image** modes.
+
+The default yaml configuration deploys AR on GPU 0 and DiT on GPU 1. You can use the default configuration file: [`glm_image.yaml`](../../../vllm_omni/model_executor/stage_configs/glm_image.yaml)
+
+| Mode           | Input        | Output | Description                        |
+| -------------- | ------------ | ------ | ---------------------------------- |
+| Text-to-Image  | Text         | Image  | Generate images from text prompts  |
+| Image-to-Image | Image + Text | Image  | Edit images with text instructions |
+
+### Text-to-Image
+
+Generate images from text prompts:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "A photorealistic mountain landscape at sunset" \
+    --height 1024 \
+    --width 1024 \
+    --output landscape.png
+```
+
+**Using curl**
+
+```bash
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [
+      {"role": "user", "content": "A beautiful sunset over the ocean with sailing boats"}
+    ],
+    "extra_body": {
+      "height": 1024,
+      "width": 1024,
+      "num_inference_steps": 50,
+      "guidance_scale": 1.5,
+      "seed": 42
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
+```
+
+Or use the script:
+
+```bash
+bash run_curl_text_to_image.sh "A futuristic city skyline at night"
+```
+
+### Image-to-Image (Image Editing)
+
+Edit images with text instructions:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "Convert this image to watercolor style" \
+    --image input.png \
+    --output watercolor.png
+```
+
+**Using curl**
+
+```bash
+IMG_B64=$(base64 < input.png | tr -d '\n')
+
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d @- <<EOF | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
+{
+  "messages": [{
+    "role": "user",
+    "content": [
+      {"type": "text", "text": "Convert this image to watercolor style"},
+      {"type": "image_url", "image_url": {"url": "data:image/png;base64,'$IMG_B64'"}}
+    ]
+  }],
+  "extra_body": {
+    "height": 1024,
+    "width": 1024,
+    "num_inference_steps": 50,
+    "guidance_scale": 1.5,
+    "seed": 42
+  }
+}
+EOF
+```
+
+Or use the script:
+
+```bash
+bash run_curl_image_edit.sh input.png "Convert to watercolor style"
+```
+
+## Generation Parameters (extra_body)
+
+| Parameter             | Type  | Default | Description                         |
+| --------------------- | ----- | ------- | ----------------------------------- |
+| `height`              | int   | 1024    | Image height in pixels              |
+| `width`               | int   | 1024    | Image width in pixels               |
+| `num_inference_steps` | int   | 50      | Number of diffusion denoising steps |
+| `guidance_scale`      | float | 1.5     | Classifier-free guidance scale      |
+| `seed`                | int   | 42      | Random seed for reproducibility     |
+| `negative_prompt`     | str   | None    | Negative prompt                     |
+
+## Response Format
+
+```json
+{
+  "id": "chatcmpl-xxx",
+  "created": 1234567890,
+  "model": "zai-org/GLM-Image",
+  "choices": [
+    {
+      "index": 0,
+      "message": {
+        "role": "assistant",
+        "content": [
+          {
+            "type": "image_url",
+            "image_url": {
+              "url": "data:image/png;base64,..."
+            }
+          }
+        ]
+      },
+      "finish_reason": "stop"
+    }
+  ],
+  "usage": {}
+}
+```
+
+## Extract Image
+
+```bash
+# From a saved JSON response
+cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
+```
+
+## Architecture
+
+GLM-Image uses a 2-stage multistage pipeline:
+
+```
+Stage 0 (AR Model)                Stage 1 (Diffusion)
+┌───────────────────┐            ┌─────────────────────┐
+│ vLLM-optimized    │  prior     │  GlmImagePipeline   │
+│ GlmImageFor       │──tokens──►│  ┌───────────────┐  │
+│ Conditional       │            │  │ DiT Denoiser  │  │
+│ Generation        │            │  └───────┬───────┘  │
+│ (9B AR model)     │            │          ▼          │
+└───────────────────┘            │  ┌───────────────┐  │
+        ▲                        │  │  VAE Decode   │──┼──► Image
+        │                        │  └───────────────┘  │
+   Text / Image                  └─────────────────────┘
+     Input
+```
+
+## VRAM Requirements
+
+| Stage             | VRAM                   |
+| :---------------- | :--------------------- |
+| Stage-0 (AR)      | **~18 GiB + KV Cache** |
+| Stage-1 (DiT+VAE) | **~20 GiB**            |
+| Total             | **~38 GiB + KV Cache** |
+
+## File Description
+
+| File                        | Description                           |
+| --------------------------- | ------------------------------------- |
+| `run_server.sh`             | Server startup script                 |
+| `run_curl_text_to_image.sh` | Text-to-image curl example            |
+| `run_curl_image_edit.sh`    | Image-to-image (editing) curl example |
+| `openai_chat_client.py`     | Python client (t2i + i2i)             |
+
+## FAQ
+
+- If you encounter OOM errors, adjust `gpu_memory_utilization` in the stage config:
+
+```yaml
+# In glm_image.yaml, reduce from default 0.6:
+gpu_memory_utilization: 0.5
+```
+
+- The first request may be slow due to model warmup. Subsequent requests will be faster.

examples/online_serving/glm_image/openai_chat_client.py

@@ -0,0 +1,172 @@
+#!/usr/bin/env python3
+"""
+GLM-Image OpenAI-compatible chat client for text-to-image and image-to-image.
+
+Usage:
+    # Text-to-image
+    python openai_chat_client.py --prompt "A cute cat" --output output.png
+
+    # Image-to-image (image editing)
+    python openai_chat_client.py --prompt "Convert to watercolor style" --image input.png --output output.png
+"""
+
+import argparse
+import base64
+from pathlib import Path
+
+import requests
+
+
+def generate_image(
+    prompt: str,
+    server_url: str = "http://localhost:8091",
+    image_path: str | None = None,
+    height: int = 1024,
+    width: int = 1024,
+    steps: int = 50,
+    guidance_scale: float = 1.5,
+    seed: int | None = None,
+    negative_prompt: str | None = None,
+) -> bytes | None:
+    """Generate or edit an image using the chat completions API.
+
+    Args:
+        prompt: Text description or editing instruction
+        server_url: Server URL
+        image_path: Path to input image (for image-to-image editing)
+        height: Image height in pixels
+        width: Image width in pixels
+        steps: Number of inference steps
+        guidance_scale: Classifier-free guidance scale
+        seed: Random seed for reproducibility
+        negative_prompt: Negative prompt
+
+    Returns:
+        Image bytes or None if failed
+    """
+    # Build message content
+    content: list[dict] = [{"type": "text", "text": prompt}]
+
+    if image_path:
+        img_path = Path(image_path)
+        if not img_path.exists():
+            print(f"Error: Image file not found: {image_path}")
+            return None
+        b64_data = base64.b64encode(img_path.read_bytes()).decode("utf-8")
+        suffix = img_path.suffix.lstrip(".").lower()
+        mime = {"jpg": "jpeg", "jpeg": "jpeg", "png": "png", "webp": "webp"}.get(suffix, "png")
+        content.append(
+            {
+                "type": "image_url",
+                "image_url": {"url": f"data:image/{mime};base64,{b64_data}"},
+            }
+        )
+
+    messages = [{"role": "user", "content": content}]
+
+    # Build request payload
+    extra_body: dict = {
+        "height": height,
+        "width": width,
+        "num_inference_steps": steps,
+        "guidance_scale": guidance_scale,
+    }
+    if seed is not None:
+        extra_body["seed"] = seed
+    if negative_prompt:
+        extra_body["negative_prompt"] = negative_prompt
+
+    payload = {"messages": messages, "extra_body": extra_body}
+
+    # Send request
+    try:
+        mode = "image-to-image" if image_path else "text-to-image"
+        print(f"Sending {mode} request to {server_url}...")
+        response = requests.post(
+            f"{server_url}/v1/chat/completions",
+            headers={"Content-Type": "application/json"},
+            json=payload,
+            timeout=600,
+        )
+        response.raise_for_status()
+        data = response.json()
+
+        # Extract image from response
+        choices = data.get("choices", [])
+        for choice in choices:
+            choice_content = choice.get("message", {}).get("content")
+            if isinstance(choice_content, list):
+                for item in choice_content:
+                    if isinstance(item, dict) and "image_url" in item:
+                        img_url = item["image_url"].get("url", "")
+                        if img_url.startswith("data:image"):
+                            _, b64 = img_url.split(",", 1)
+                            return base64.b64decode(b64)
+
+        print(f"Unexpected response format: {data}")
+        return None
+
+    except Exception as e:
+        print(f"Error: {e}")
+        return None
+
+
+def main():
+    parser = argparse.ArgumentParser(description="GLM-Image chat client")
+    parser.add_argument(
+        "--prompt",
+        "-p",
+        default="A beautiful sunset over the ocean with sailing boats",
+        help="Text prompt",
+    )
+    parser.add_argument("--output", "-o", default="glm_image_output.png", help="Output file")
+    parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
+
+    # Image-to-image
+    parser.add_argument(
+        "--image",
+        "-i",
+        type=str,
+        help="Input image path (for image-to-image editing)",
+    )
+
+    # Generation parameters
+    parser.add_argument("--height", type=int, default=1024, help="Image height")
+    parser.add_argument("--width", type=int, default=1024, help="Image width")
+    parser.add_argument("--steps", type=int, default=50, help="Inference steps")
+    parser.add_argument("--guidance-scale", type=float, default=1.5, help="CFG guidance scale")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    parser.add_argument("--negative", help="Negative prompt")
+
+    args = parser.parse_args()
+
+    mode = "image-to-image" if args.image else "text-to-image"
+    print(f"Mode: {mode}")
+    print(f"Prompt: {args.prompt}")
+    if args.image:
+        print(f"Input image: {args.image}")
+
+    image_bytes = generate_image(
+        prompt=args.prompt,
+        server_url=args.server,
+        image_path=args.image,
+        height=args.height,
+        width=args.width,
+        steps=args.steps,
+        guidance_scale=args.guidance_scale,
+        seed=args.seed,
+        negative_prompt=args.negative,
+    )
+
+    if image_bytes:
+        output_path = Path(args.output)
+        output_path.write_bytes(image_bytes)
+        print(f"Image saved to: {output_path}")
+        print(f"Size: {len(image_bytes) / 1024:.1f} KB")
+    else:
+        print("Failed to generate image")
+        exit(1)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/glm_image/run_curl_image_edit.sh

@@ -0,0 +1,61 @@
+#!/bin/bash
+# GLM-Image image-edit (image-to-image) curl example
+
+set -euo pipefail
+
+if [[ $# -lt 2 ]]; then
+  echo "Usage: $0 <input_image> \"<edit_prompt>\" [output_file]" >&2
+  exit 1
+fi
+
+INPUT_IMG=$1
+PROMPT=$2
+SERVER="${SERVER:-http://localhost:8091}"
+CURRENT_TIME=$(date +%Y%m%d%H%M%S)
+OUTPUT="${3:-glm_image_i2i_${CURRENT_TIME}.png}"
+
+if [[ ! -f "$INPUT_IMG" ]]; then
+  echo "Input image not found: $INPUT_IMG" >&2
+  exit 1
+fi
+
+# base64 encode (macOS uses -i, Linux uses -w0)
+if [[ "$(uname)" == "Darwin" ]]; then
+  IMG_B64=$(base64 < "$INPUT_IMG" | tr -d '\n')
+else
+  IMG_B64=$(base64 -w0 "$INPUT_IMG")
+fi
+
+REQUEST_JSON=$(
+  jq -n --arg prompt "$PROMPT" --arg img "$IMG_B64" '{
+    messages: [{
+      role: "user",
+      content: [
+        {"type": "text", "text": $prompt},
+        {"type": "image_url", "image_url": {"url": ("data:image/png;base64," + $img)}}
+      ]
+    }],
+    extra_body: {
+      height: 1024,
+      width: 1024,
+      num_inference_steps: 50,
+      guidance_scale: 1.5,
+      seed: 42
+    }
+  }'
+)
+
+echo "Generating edited image..."
+echo "Server: $SERVER"
+echo "Prompt: $PROMPT"
+echo "Input : $INPUT_IMG"
+echo "Output: $OUTPUT"
+
+curl -s "$SERVER/v1/chat/completions" \
+  -H "Content-Type: application/json" \
+  -d "$REQUEST_JSON" \
+  | jq -r '.choices[0].message.content[0].image_url.url' \
+  | cut -d',' -f2- \
+  | base64 -d > "$OUTPUT"
+
+echo "Image saved to: $OUTPUT"

examples/online_serving/glm_image/run_curl_text_to_image.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+# GLM-Image text-to-image curl example
+
+set -euo pipefail
+
+PROMPT="${1:-A beautiful sunset over the ocean with sailing boats}"
+SERVER="${SERVER:-http://localhost:8091}"
+OUTPUT="${OUTPUT:-glm_image_t2i_output.png}"
+
+echo "Generating image..."
+echo "Server: $SERVER"
+echo "Prompt: $PROMPT"
+echo "Output: $OUTPUT"
+
+curl -s "$SERVER/v1/chat/completions" \
+  -H "Content-Type: application/json" \
+  -d "{
+    \"messages\": [
+      {\"role\": \"user\", \"content\": \"$PROMPT\"}
+    ],
+    \"extra_body\": {
+      \"height\": 1024,
+      \"width\": 1024,
+      \"num_inference_steps\": 50,
+      \"guidance_scale\": 1.5,
+      \"seed\": 42
+    }
+  }" | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > "$OUTPUT"
+
+echo "Image saved to: $OUTPUT"

examples/online_serving/glm_image/run_server.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+# GLM-Image online serving startup script
+
+MODEL="${MODEL:-zai-org/GLM-Image}"
+PORT="${PORT:-8091}"
+
+echo "Starting GLM-Image server..."
+echo "Model: $MODEL"
+echo "Port: $PORT"
+
+vllm serve "$MODEL" --omni \
+    --port "$PORT"

examples/online_serving/mimo_audio/README.md

@@ -0,0 +1,58 @@
+# Online serving Example of vLLM-Omni for MiMo-Audio
+
+## 🛠️ Installation
+
+Please refer to [README.md](../../../README.md)
+
+## Run examples (MiMo-Audio)
+
+### Launch the Server
+```bash
+export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
+
+vllm-omni serve XiaomiMiMo/MiMo-Audio-7B-Instruct --omni \
+--served-model-name "MiMo-Audio-7B-Instruct"  \
+--port 18091 --stage-configs-path ./vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
+--chat-template ./examples/online_serving/mimo_audio/chat_template.jinja
+```
+> ⚠️ **Important**  
+> **MiMo-Audio is not compatible with the default chat template.**  
+> The provided `chat_template.jinja` implements MiMo-specific role, audio token, and instruction formatting and **must be used for all inference**.
+
+
+### Send Multi-modal Request
+
+Get into the example folder
+```bash
+cd examples/online_serving/mimo_audio
+```
+
+####  Send request via python
+
+```bash
+# Audio dialogue task
+python openai_chat_completion_client_for_multimodal_generation.py \
+--query-type multi_audios \
+--message-json ../../offline_inference/mimo_audio/message_base64_wav.json
+```
+
+The Python client supports the following command-line arguments:
+
+- `--query-type` (or `-q`): Query type (default: `multi_audios`)
+  - Options: `multi_audios`, `text`
+- `--message-json` (or `-m`): Path to `base64` multi rounds audio messages json file
+  - Do not pass any value for "text" query type
+  - Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs, only for "Are these two audio clips the same?" task
+  - Example: `---message-json ./examples/offline_inference/mimo_audio/message_base64_wav.json`
+- `--prompt` (or `-p`): Custom text prompt/question, only for query type is "text"(TTS task)
+  - Attention! Do not pass any value for "multi_audios" query type
+  - Example: `--prompt "What are the main activities shown in this video?"`
+
+
+For example, to use multi rounds audios with local files:
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+--query-type multi_audios \
+--message-json ../../offline_inference/mimo_audio/message_base64_wav.json
+```

examples/online_serving/mimo_audio/chat_template.jinja

@@ -0,0 +1,86 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- messages[0]['content'] }}
+    {%- else %}
+        {{- 'You are a helpful assistant.' }}
+    {%- endif %}
+    {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0]['role'] == 'system' %}
+        {%- set _m = '<|sosp|><|empty|><|eosp|>' -%}
+        {%- set _raw0 = messages[0]['content'] if messages[0]['content'] is string else '' -%}
+        {%- if _m in _raw0 %}
+            {%- set _t0 = (_raw0 | replace(_m ~ '\n', '') | replace(_m, '') | trim) -%}
+            {{- '<|im_start|>system\n' + (_t0 ~ _m if _t0 else _m) + '<|im_end|>\n' }}
+        {%- else %}
+            {{- '<|im_start|>system\n' + _raw0 + '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+
+{%- for message in messages %}
+    {%- if message['role'] == 'assistant' %}
+        {{- '<|im_start|>assistant' }}
+        {%- set _sosp = '<|sosp|><|empty|><|eosp|>' -%}
+        {%- set _text = message['content'] if message['content'] is string else '' -%}
+        {%- if _sosp in _text %}
+            {%- set _clean = _text | replace(_sosp, '') -%}
+            {%- set _body = _clean[1:] if (_clean and _clean[0] == '\n') else _clean -%}
+            {{- '\n<|sostm|>' + _body + '<|eot|><|empty|><|eostm|>' }}
+        {%- else %}
+            {%- set _body = _text[1:] if (_text and _text[0] == '\n') else _text -%}
+            {{- '\n<|sostm|>' + _body + '<|eot|><|eostm|>' }}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+
+    {%- elif message['role'] == 'user' %}
+        {%- set _m = '<|sosp|><|empty|><|eosp|>' -%}
+        {%- set _raw = message['content'] if message['content'] is string else '' -%}
+        {%- if _m in _raw %}
+            {%- set _t = (_raw | replace(_m ~ '\n', '') | replace(_m, '') | trim) -%}
+            {{- '<|im_start|>user\n' + (_t ~ _m if _t else _m) + '<|im_end|>\n' }}
+        {%- else %}
+            {{- '<|im_start|>user\n' + _raw + '<|im_end|>\n' }}
+        {%- endif %}
+
+    {%- elif message['role'] == 'system' %}
+        {%- if not loop.first %}
+            {%- set _m = '<|sosp|><|empty|><|eosp|>' -%}
+            {%- set _raw = message['content'] if message['content'] is string else '' -%}
+            {%- if _m in _raw %}
+                {%- set _t = (_raw | replace(_m ~ '\n', '') | replace(_m, '') | trim) -%}
+                {{- '<|im_start|>system\n' + (_t ~ _m if _t else _m) + '<|im_end|>\n' }}
+            {%- else %}
+                {{- '<|im_start|>system\n' + _raw + '<|im_end|>\n' }}
+            {%- endif %}
+        {%- endif %}
+
+    {%- elif message['role'] == 'tool' %}
+        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1]['role'] != 'tool') %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {%- set _m = '<|sosp|><|empty|><|eosp|>' -%}
+        {%- set _raw = message['content'] if message['content'] is string else '' -%}
+        {%- if _m in _raw %}
+            {%- set _t = (_raw | replace(_m ~ '\n', '') | replace(_m, '') | trim) -%}
+            {{- '\n<tool_response>\n' + (_t ~ _m if _t else _m) + '\n</tool_response>' }}
+        {%- else %}
+            {{- '\n<tool_response>\n' + _raw + '\n</tool_response>' }}
+        {%- endif %}
+        {%- if loop.last or (messages[loop.index0 + 1]['role'] != 'tool') %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n<|sostm|>' }}
+{%- endif %}

examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py

@@ -0,0 +1,545 @@
+import base64
+import concurrent.futures
+import datetime
+import json
+import os
+import queue
+import sys
+import threading
+import time
+from typing import Any
+
+import requests
+from openai import OpenAI
+from vllm.assets.audio import AudioAsset
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+# Modify OpenAI's API key and API base to use vLLM's API server.
+openai_api_key = "EMPTY"
+openai_api_base = "http://localhost:18091/v1"
+
+
+# Modify OpenAI's API key and API base to use vLLM's API server.
+
+
+client = OpenAI(
+    # defaults to os.environ.get("OPENAI_API_KEY")
+    api_key=openai_api_key,
+    base_url=openai_api_base,
+)
+
+SEED = 42
+
+
+def encode_base64_content_from_url(content_url: str) -> str:
+    """Encode a content retrieved from a remote url to base64 format."""
+
+    with requests.get(content_url) as response:
+        response.raise_for_status()
+        result = base64.b64encode(response.content).decode("utf-8")
+
+    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: str | None) -> 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: str | None) -> 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: str | None) -> 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 already a base64 data URL (starts with "data:"), returns it as-is.
+    If audio_path is a URL (starts with http:// or https://), returns it as-is.
+    If audio_path is a local file path, encodes it to base64 data URL.
+    """
+    if not audio_path:
+        # Default audio URL
+        return AudioAsset("mary_had_lamb").url
+
+    # Check if it's already a base64 data URL
+    if audio_path.startswith("data:"):
+        return audio_path
+
+    # 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 load_messages_from_json(message_json_path: str) -> dict[str, Any]:
+    """Load messages from a JSON file."""
+    with open(message_json_path, encoding="utf-8") as f:
+        data = json.load(f)
+    return data
+
+
+def process_audio_url_in_content(content: dict[str, Any]) -> dict[str, Any]:
+    """Process audio_url in content, handling both file paths and base64 URLs."""
+    if content.get("type") == "audio_url":
+        audio_url = content.get("audio_url", {}).get("url")
+        if audio_url:
+            # Process the audio URL (handles base64, file paths, and URLs)
+            processed_url = get_audio_url_from_path(audio_url)
+            content = content.copy()
+            content["audio_url"] = {"url": processed_url}
+    return content
+
+
+def get_system_prompt(message_json_path: str | None = None):
+    """Get system prompt, optionally from message.json file."""
+    if message_json_path and os.path.exists(message_json_path):
+        data = load_messages_from_json(message_json_path)
+        messages = data.get("messages", [])
+        # Find the first system message
+        for msg in messages:
+            if msg.get("role") == "system":
+                # Process audio URLs in the content
+                processed_content = []
+                for content_item in msg.get("content", []):
+                    processed_item = process_audio_url_in_content(content_item)
+                    processed_content.append(processed_item)
+                return {
+                    "role": "system",
+                    "content": processed_content,
+                }
+
+
+def get_text_query(custom_prompt: str | None = None):
+    question = f"Please convert this text to speech: {custom_prompt}"
+    prompt = {
+        "role": "user",
+        "content": [
+            {
+                "type": "text",
+                "text": f"{question}",
+            }
+        ],
+    }
+    return prompt
+
+
+def get_multi_audios_query(
+    audio_path: str | None = None,
+    custom_prompt: str | None = None,
+    message_json_path: str | None = None,
+) -> list[dict[str, Any]] | dict[str, Any]:
+    """Get multi-audios query, optionally from message.json file."""
+    if message_json_path and os.path.exists(message_json_path):
+        data = load_messages_from_json(message_json_path)
+        messages = data.get("messages", [])
+
+        # Skip the first system role and return all other messages
+        rebuild_prompt_messages = []
+        skipped_first_system = False
+
+        for msg in messages:
+            if msg.get("role") == "system" and not skipped_first_system:
+                skipped_first_system = True
+                continue
+
+            # Process all content items in the message
+            processed_content = []
+            for content_item in msg.get("content", []):
+                processed_item = process_audio_url_in_content(content_item)
+                processed_content.append(processed_item)
+
+            rebuild_prompt_messages.append(
+                {
+                    "role": msg.get("role"),
+                    "content": processed_content,
+                }
+            )
+
+        return rebuild_prompt_messages
+
+    # Original behavior when message_json_path is not provided
+    question = custom_prompt or "Are these two audio clips the same?"
+    audio_url = get_audio_url_from_path(audio_path)
+    prompt = {
+        "role": "user",
+        "content": [
+            {
+                "type": "audio_url",
+                "audio_url": {"url": audio_url},
+            },
+            {
+                "type": "audio_url",
+                "audio_url": {"url": AudioAsset("winning_call").url},
+            },
+            {
+                "type": "text",
+                "text": f"{question}",
+            },
+        ],
+    }
+    return prompt
+
+
+query_map = {
+    "multi_audios": get_multi_audios_query,
+    "text": get_text_query,
+}
+
+
+def run_multimodal_generation(args) -> None:
+    model_name = "MiMo-Audio-7B-Instruct"
+    thinker_sampling_params = {
+        "temperature": 0.0,  # Deterministic - no randomness
+        "top_p": 1.0,  # Disable nucleus sampling
+        "top_k": -1,  # Disable top-k sampling
+        "max_tokens": 3200,
+        "seed": SEED,  # Fixed seed for sampling
+        "detokenize": True,
+        "repetition_penalty": 1.1,
+    }
+    code2wav_sampling_params = {
+        "temperature": 0.0,  # Deterministic - no randomness
+        "top_p": 1.0,  # Disable nucleus sampling
+        "top_k": -1,  # Disable top-k sampling
+        "max_tokens": 4096 * 16,
+        "seed": SEED,  # Fixed seed for sampling
+        "detokenize": True,
+        "repetition_penalty": 1.1,
+    }
+
+    sampling_params_list = [
+        thinker_sampling_params,
+        code2wav_sampling_params,
+    ]
+
+    # Get paths and custom prompt from args
+    audio_path = getattr(args, "audio_path", None)
+    custom_prompt = getattr(args, "prompt", None)
+    message_json_path = getattr(args, "message_json", None)
+    output_audio_path = getattr(args, "output_audio_path", None)
+
+    # Get the query function and call it with appropriate parameters
+    query_func = query_map[args.query_type]
+    if args.query_type == "multi_audios":
+        prompt = query_func(audio_path=audio_path, custom_prompt=custom_prompt, message_json_path=message_json_path)
+    elif args.query_type == "text":
+        prompt = query_func(custom_prompt=custom_prompt)
+    else:
+        prompt = query_func()
+
+    extra_body = {
+        "sampling_params_list": sampling_params_list
+        # Optional, it has a default setting in stage_configs of the corresponding model.
+    }
+
+    # Build messages list (filter None so concurrent tasks get valid structure)
+    system = get_system_prompt(message_json_path=message_json_path)
+    if args.query_type == "multi_audios" and isinstance(prompt, list):
+        messages = ([system] if system else []) + prompt
+    elif args.query_type == "text":
+        messages = [prompt]
+
+    num_concurrent_requests = args.num_concurrent_requests
+    start_time = time.perf_counter()
+
+    def run_one_request(req_id: int):
+        """Submit one request and return (req_id, response) to preserve request order."""
+        return req_id, client.chat.completions.create(
+            messages=messages,
+            model=model_name,
+            extra_body=extra_body,
+            stream=args.stream,
+        )
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
+        futures = [executor.submit(run_one_request, req_id) for req_id in range(num_concurrent_requests)]
+
+        if not args.stream:
+            # Collect by req_id so chat_completions[i] = response for request i
+            results_by_req = [None] * num_concurrent_requests
+            for future in concurrent.futures.as_completed(futures):
+                req_id, completion = future.result()
+                results_by_req[req_id] = completion
+            chat_completions = results_by_req
+
+            assert len(chat_completions) == num_concurrent_requests
+            count = 0
+            for req_id, chat_completion in enumerate(chat_completions):
+                chat_completion_id = getattr(chat_completion, "id", "")
+                for choice in chat_completion.choices:
+                    if choice.message.audio:
+                        audio_data = base64.b64decode(choice.message.audio.data)
+                        audio_file_path = f"{output_audio_path}/{args.query_type}/audio_{count}.wav"
+                        os.makedirs(os.path.dirname(audio_file_path), exist_ok=True)
+                        with open(audio_file_path, "wb") as f:
+                            f.write(audio_data)
+                        print(f"[req {req_id}_{chat_completion_id}] Audio saved to {audio_file_path}")
+                        count += 1
+                    elif choice.message.content:
+                        print(
+                            f"[req {req_id}_{chat_completion_id}] Chat completion output from text:",
+                            choice.message.content,
+                        )
+        else:
+            # Stream mode: process chunks from all requests in real-time,
+            # displaying them in the order they arrive.
+            chunk_queue = queue.Queue()
+            audio_counters = {}
+            chat_completion_id_by_req = {}
+            chat_completion_id_lock = threading.Lock()
+            last_text_req_id = None
+
+            def _stream_reader(req_id: int, stream):
+                """Read one stream and relay every chunk to the shared queue."""
+                try:
+                    for chunk in stream:
+                        chunk_queue.put((req_id, time.perf_counter(), chunk))
+                    chunk_queue.put((req_id, time.perf_counter(), None))
+                except Exception as e:
+                    print(f"\n[Request {req_id}] stream error: {e}", file=sys.stderr, flush=True)
+                    chunk_queue.put((req_id, time.time(), None))
+
+            # Kick off a reader thread per request
+            reader_threads = []
+            for req_id, future in enumerate(futures):
+                req_id_from_future, stream = future.result()
+                assert req_id == req_id_from_future, f"Request ID mismatch: {req_id} != {req_id_from_future}"
+                audio_counters[req_id] = 0
+                t = threading.Thread(target=_stream_reader, args=(req_id, stream), daemon=True)
+                t.start()
+                reader_threads.append(t)
+
+            # Main loop – consume chunks in arrival order
+            active_streams = set(range(num_concurrent_requests))
+            while active_streams:
+                try:
+                    request_id, ts, chunk = chunk_queue.get(timeout=2.0)
+                except queue.Empty:
+                    if all(not t.is_alive() for t in reader_threads):
+                        break
+                    continue
+
+                if chunk is None:
+                    elapsed = ts - start_time
+                    print(f"  ({elapsed:.2f}s)", flush=True)
+                    with chat_completion_id_lock:
+                        chat_completion_id = chat_completion_id_by_req.get(request_id, "")
+                    print(
+                        f" [req {request_id}_{chat_completion_id}] Time finished for streaming: ",
+                        datetime.datetime.now(),
+                        file=sys.stderr,
+                        flush=True,
+                    )
+                    active_streams.discard(request_id)
+                    continue
+
+                with chat_completion_id_lock:
+                    if request_id not in chat_completion_id_by_req:
+                        chat_completion_id_by_req[request_id] = getattr(chunk, "id", "")
+                    chat_completion_id = chat_completion_id_by_req[request_id]
+
+                modality = getattr(chunk, "modality", None)
+                elapsed = ts - start_time
+                for choice in chunk.choices:
+                    content = getattr(choice.delta, "content", None) if hasattr(choice, "delta") else None
+
+                    if modality == "audio" and content:
+                        audio_data = base64.b64decode(content)
+                        audio_dir = (
+                            f"{output_audio_path}/{args.query_type}/{num_concurrent_requests}/{chat_completion_id}"
+                        )
+                        os.makedirs(audio_dir, exist_ok=True)
+                        audio_file_path = f"{audio_dir}/audio_{audio_counters[request_id]}.wav"
+                        with open(audio_file_path, "wb") as f:
+                            f.write(audio_data)
+                        print(
+                            f"\n[{elapsed:7.2f}s][req {request_id}_{chat_completion_id}] Audio saved to {audio_file_path}",
+                            file=sys.stderr,
+                            flush=True,
+                        )
+                        audio_counters[request_id] += 1
+
+                    elif modality == "text" and content:
+                        if last_text_req_id != request_id:
+                            if last_text_req_id is not None:
+                                print(flush=True)
+                            print(
+                                f"\n[{elapsed:7.2f}s][req {request_id}_{chat_completion_id}] content:",
+                                end="",
+                                flush=True,
+                            )
+                            last_text_req_id = request_id
+                        print(
+                            f"\n[{elapsed:7.2f}s][req {request_id}_{chat_completion_id}] {content}", end="", flush=True
+                        )
+
+            # Final newline if the last output was streaming text
+            if last_text_req_id is not None:
+                print(flush=True)
+
+            for t in reader_threads:
+                t.join(timeout=1.0)
+
+            elapsed = time.perf_counter() - start_time
+            print(
+                f"num_concurrent_requests_{num_concurrent_requests}>>>>>>>>>Time finished for streaming<<<<<<<<: ",
+                elapsed,
+                file=sys.stderr,
+                flush=True,
+            )
+            timing_audio_folder = f"{output_audio_path}/{args.query_type}/{num_concurrent_requests}"
+            os.makedirs(timing_audio_folder, exist_ok=True)
+            timing_file = os.path.join(timing_audio_folder, "streaming_finish_time.txt")
+            with open(timing_file, "w") as f:
+                f.write(f"num_concurrent_requests_{num_concurrent_requests} elapsed_seconds: {elapsed}\n")
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="multi_audios",
+        choices=query_map.keys(),
+        help="Query type.",
+    )
+    parser.add_argument(
+        "--prompt",
+        "-p",
+        type=str,
+        default="Who are you?",
+        help="Custom text prompt/question to use instead of the default prompt for the selected query type.",
+    )
+    parser.add_argument(
+        "--message-json",
+        "-m",
+        type=str,
+        # default="../../offline_inference/mimo_audio/message_base64_wav.json",
+        default="../../offline_inference/mimo_audio/message_base64_wav_tts.json",
+        help="Path to message.json file containing conversation history. When provided, "
+        "system prompt and multi_audios query will be loaded from this file.",
+    )
+    parser.add_argument(
+        "--output-audio-path",
+        "-o",
+        type=str,
+        default="./",
+        help="Path to save the generated audio files.",
+    )
+    parser.add_argument(
+        "--stream",
+        action="store_true",
+        help="Stream the response.",
+    )
+    parser.add_argument(
+        "--num-concurrent-requests",
+        type=int,
+        default=1,
+        help="Number of concurrent requests to send. Default is 1.",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_multimodal_generation(args)

examples/online_serving/qwen3_tts/README.md

@@ -1,40 +1,75 @@
 # Qwen3-TTS
 
-This directory contains examples for running Qwen3-TTS models with vLLM-Omni's online serving API.
+## 🛠️ Installation
+
+Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
 
 ## Supported Models
 
-| Model | Task Type | Description |
-|-------|-----------|-------------|
+| Model                                  | Task Type   | Description                                           |
+| -------------------------------------- | ----------- | ----------------------------------------------------- |
 | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
-| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
-| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
-| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
-| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description              |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-Base`        | Base        | Voice cloning from reference audio                    |
+| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant                                |
+| `Qwen/Qwen3-TTS-12Hz-0.6B-Base`        | Base        | Smaller/faster variant for voice cloning              |
 
-## Quick Start
+## Run examples (Qwen3-TTS)
 
-### 1. Start the Server
-
-```bash
-# CustomVoice model (default)
-./run_server.sh
-
-# Or specify task type
-./run_server.sh CustomVoice
-./run_server.sh VoiceDesign
-./run_server.sh Base
-```
-
-Or launch directly with vllm serve:
+### Launch the Server
 
 ```bash
+# CustomVoice model (predefined speakers)
 vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
     --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
-    --omni --port 8091 --trust-remote-code --enforce-eager
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager
+
+# VoiceDesign model
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager
+
+# Base model (voice cloning)
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager
 ```
 
-### 2. Run the Client
+If you have custom stage configs file, launch the server with command below
+```bash
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+    --stage-configs-path /path/to/stage_configs_file \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager
+```
+
+Alternatively, use the convenience script:
+```bash
+./run_server.sh                  # Default: CustomVoice model
+./run_server.sh CustomVoice      # CustomVoice model
+./run_server.sh VoiceDesign      # VoiceDesign model
+./run_server.sh Base             # Base (voice clone) model
+```
+
+### Send TTS Request
+
+Get into the example folder
+```bash
+cd examples/online_serving/qwen3_tts
+```
+
+####  Send request via python
 
 ```bash
 # CustomVoice: Use predefined speaker
@@ -65,7 +100,21 @@ python openai_speech_client.py \
     --ref-text "Original transcript of the reference audio"
 ```
 
-### 3. Using curl
+The Python client supports the following command-line arguments:
+
+- `--api-base`: API base URL (default: `http://localhost:8091`)
+- `--model` (or `-m`): Model name/path (default: `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice`)
+- `--task-type` (or `-t`): TTS task type. Options: `CustomVoice`, `VoiceDesign`, `Base`
+- `--text`: Text to synthesize (required)
+- `--voice`: Speaker/voice name (default: `vivian`). Options: `vivian`, `ryan`, `aiden`, etc.
+- `--language`: Language. Options: `Auto`, `Chinese`, `English`, `Japanese`, `Korean`, `German`, `French`, `Russian`, `Portuguese`, `Spanish`, `Italian`
+- `--instructions`: Voice style/emotion instructions
+- `--ref-audio`: Reference audio file path or URL for voice cloning (Base task)
+- `--ref-text`: Reference audio transcript for voice cloning (Base task)
+- `--response-format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`, `aac`, `opus`
+- `--output` (or `-o`): Output audio file path (default: `tts_output.wav`)
+
+####  Send request via curl
 
 ```bash
 # Simple TTS request
@@ -90,24 +139,48 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 curl http://localhost:8091/v1/audio/voices
 ```
 
-## Server Configuration
+### Using OpenAI SDK
 
-### TTS-Specific CLI Parameters
+```python
+from openai import OpenAI
 
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `--tts-max-instructions-length` | int | 500 | Maximum character length for voice style instructions |
+client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
 
-**Example with custom instructions length:**
+response = client.audio.speech.create(
+    model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+    voice="vivian",
+    input="Hello, how are you?",
+)
 
-```bash
-vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
-    --omni \
-    --port 8000 \
-    --tts-max-instructions-length 1000
+response.stream_to_file("output.wav")
 ```
 
-This is useful when you need longer voice style descriptions, especially for the VoiceDesign task type.
+### Using Python httpx
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8091/v1/audio/speech",
+    json={
+        "input": "Hello, how are you?",
+        "voice": "vivian",
+        "language": "English",
+    },
+    timeout=300.0,
+)
+
+with open("output.wav", "wb") as f:
+    f.write(response.content)
+```
+
+### FAQ
+
+If you encounter error about backend of librosa, try to install ffmpeg with command below.
+```
+sudo apt update
+sudo apt install ffmpeg
+```
 
 ## API Reference
 
@@ -115,6 +188,7 @@ This is useful when you need longer voice style descriptions, especially for the
 
 ```
 POST /v1/audio/speech
+Content-Type: application/json
 ```
 
 This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/docs/api-reference/audio/createSpeech) format with additional Qwen3-TTS parameters.
@@ -154,84 +228,69 @@ Lists available voices for the loaded model:
 
 ### Response
 
-Returns audio data in the requested format (default: WAV).
+Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
 
 ## Parameters
 
 ### OpenAI Standard Parameters
 
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `input` | string | **required** | Text to synthesize |
-| `model` | string | server's model | Model to use (optional, should match server if specified) |
-| `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
-| `speed` | float | 1.0 | Playback speed (0.25-4.0) |
+| Parameter         | Type   | Default        | Description                                                 |
+| ----------------- | ------ | -------------- | ----------------------------------------------------------- |
+| `input`           | string | **required**   | Text to synthesize                                          |
+| `model`           | string | server's model | Model to use (optional, should match server if specified)   |
+| `voice`           | string | "vivian"       | Speaker name (e.g., vivian, ryan, aiden)                    |
+| `response_format` | string | "wav"          | Audio format: wav, mp3, flac, pcm, aac, opus                |
+| `speed`           | float  | 1.0            | Playback speed (0.25-4.0, not supported with `stream=true`) |
 
 ### vLLM-Omni Extension Parameters
 
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `task_type` | string | "CustomVoice" | Task: CustomVoice, VoiceDesign, or Base |
-| `language` | string | "Auto" | Language (see supported languages below) |
-| `instructions` | string | "" | Voice style/emotion instructions |
-| `max_new_tokens` | int | 2048 | Maximum tokens to generate |
+| Parameter        | Type   | Default       | Description                                                                  |
+| ---------------- | ------ | ------------- | ---------------------------------------------------------------------------- |
+| `task_type`      | string | "CustomVoice" | Task: CustomVoice, VoiceDesign, or Base                                      |
+| `language`       | string | "Auto"        | Language (see supported languages below)                                     |
+| `instructions`   | string | ""            | Voice style/emotion instructions                                             |
+| `max_new_tokens` | int    | 2048          | Maximum tokens to generate                                                   |
+| `stream`         | bool   | false         | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
 
 **Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
 
 ### Voice Clone Parameters (Base task)
 
-| Parameter | Type | Default | Description |
-|-----------|------|----------|-------------|
-| `ref_audio` | string | null | Reference audio (URL or base64 data URL) |
-| `ref_text` | string | null | Transcript of reference audio |
-| `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) |
+| Parameter            | Type   | Required | Description                                  |
+| -------------------- | ------ | -------- | -------------------------------------------- |
+| `ref_audio`          | string | **Yes**  | Reference audio (URL or base64 data URL)     |
+| `ref_text`           | string | No       | Transcript of reference audio (for ICL mode) |
+| `x_vector_only_mode` | bool   | No       | Use speaker embedding only (no ICL)          |
 
-## Python Usage
+## Streaming
 
-### Using OpenAI SDK
+Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded
+(one chunk per 25-frame Code2Wav window):
 
-```python
-from openai import OpenAI
-
-client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
-
-response = client.audio.speech.create(
-    model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
-    voice="vivian",
-    input="Hello, how are you?",
-)
-
-response.stream_to_file("output.wav")
-```
-
-### Using httpx
-
-```python
-import httpx
-
-response = httpx.post(
-    "http://localhost:8091/v1/audio/speech",
-    json={
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
         "input": "Hello, how are you?",
         "voice": "vivian",
         "language": "English",
-    },
-    timeout=300.0,
-)
-
-with open("output.wav", "wb") as f:
-    f.write(response.content)
+        "stream": true,
+        "response_format": "pcm"
+    }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
 ```
 
+**Constraints:**
+- `stream=true` requires `response_format="pcm"` (raw 16-bit signed PCM, 24 kHz mono).
+- `speed` adjustment is not supported when streaming.
+- Requires the server stage config to have `async_chunk: true` (default in `qwen3_tts.yaml`).
+
 ## Limitations
 
-- **No streaming**: Audio is generated completely before being returned. Streaming will be supported after the pipeline is disaggregated (see RFC #938).
 - **Single request**: Batch processing is not yet optimized for online serving.
 
 ## Troubleshooting
 
-1. **"TTS model did not produce audio output"**: Ensure you're using the correct model variant for your task type (CustomVoice task → CustomVoice model, etc.)
+1. **TTS model did not produce audio output**: Ensure you're using the correct model variant for your task type (CustomVoice task → CustomVoice model, etc.)
 2. **Connection refused**: Make sure the server is running on the correct port
 3. **Out of memory**: Use smaller model variant (`Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice`) or reduce `--gpu-memory-utilization`
 4. **Unsupported speaker**: Use `/v1/audio/voices` to list available voices for the loaded model

pyproject.toml

@@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "vllm-omni"
-version = "0.14.0"
+dynamic = ["version", "dependencies"]
 description = "A framework for efficient model inference with omni-modality models"
 readme = "README.md"
 requires-python = ">=3.10,<3.14"
@@ -17,7 +17,6 @@ authors = [
     {name = "vLLM-Omni Team"}
 ]
 keywords = ["vllm", "multimodal", "diffusion", "transformer", "inference", "serving"]
-dynamic = ["dependencies"]
 classifiers = [
     "Development Status :: 3 - Alpha",
     "Intended Audience :: Developers",
@@ -49,7 +48,8 @@ dev = [
     "soundfile>=0.13.1",
     "imageio[ffmpeg]>=0.6.0",
     "opencv-python>=4.12.0.88",
-    "mooncake-transfer-engine==0.3.8.post1"
+    "mooncake-transfer-engine==0.3.8.post1",
+    "av" # for ComfyUI tests
 ]
 
 docs = [
@@ -84,8 +84,12 @@ where = ["."]
 include = ["vllm_omni*"]
 
 [tool.setuptools.package-data]
+"vllm_omni" = ["_version.py", "py.typed"]
 "vllm_omni.model_executor.stage_configs" = ["*.yaml"]
 
+[tool.setuptools_scm]
+# no extra settings needed, presence enables setuptools-scm
+
 [tool.ruff]
 line-length = 120
 exclude = [

setup.py

@@ -12,6 +12,7 @@ import sys
 from pathlib import Path
 
 from setuptools import setup
+from setuptools_scm import get_version
 
 
 def uninstall_onnxruntime() -> None:
@@ -104,6 +105,66 @@ def detect_target_device() -> str:
         return "cuda"
 
 
+def get_vllm_omni_version() -> str:
+    """
+    Get the vLLM-Omni version with device-specific suffix.
+
+    Version format: {base_version}+{device}
+    Examples:
+        - 0.14.0+cuda (release version with CUDA)
+        - 0.14.1.dev23+g1a2b3c4.rocm (dev version with ROCm)
+        - 0.15.0+npu (release version with NPU)
+
+    Environment variables:
+        VLLM_OMNI_VERSION_OVERRIDE: Override version completely
+        VLLM_OMNI_TARGET_DEVICE: Override device detection
+
+    Returns:
+        Version string with device suffix
+    """
+    # Allow complete version override via environment variable
+    if env_version := os.getenv("VLLM_OMNI_VERSION_OVERRIDE"):
+        print(f"Overriding vLLM-Omni version with {env_version} from VLLM_OMNI_VERSION_OVERRIDE")
+        os.environ["SETUPTOOLS_SCM_PRETEND_VERSION"] = env_version
+        # Get version without device suffix for override case
+        version = get_version(write_to="vllm_omni/_version.py")
+    else:
+        # Generate version from git tags via setuptools_scm (without writing yet)
+        try:
+            version = get_version()
+        except Exception as e:
+            print(f"Warning: Failed to get version from git, using fallback: {e}")
+            version = "dev"
+
+        # Determine separator: '+' for normal versions, '.' for dev versions with '+'
+        sep = "+" if "+" not in version else "."
+
+        # Append device-specific suffix
+        device = detect_target_device()
+
+        if device == "cuda":
+            # if it is cuda, following vLLM
+            # we don't need to add any suffix
+            pass
+        elif device == "rocm":
+            version += f"{sep}rocm"
+        elif device == "npu":
+            version += f"{sep}npu"
+        elif device == "xpu":
+            version += f"{sep}xpu"
+        elif device == "cpu":
+            version += f"{sep}cpu"
+        else:
+            raise RuntimeError(f"Unknown target device: {device}")
+
+    # Tell setuptools_scm to use this version
+    # This will be picked up by write_to and written to _version.py
+    os.environ["SETUPTOOLS_SCM_PRETEND_VERSION"] = version
+    get_version(write_to="vllm_omni/_version.py")  # Write the correct version
+    print(f"Generated version: {version}")
+    return version
+
+
 def load_requirements(file_path: Path) -> list[str]:
     """
     Load requirements from a file, supporting -r directive for recursive loading.
@@ -167,5 +228,6 @@ if __name__ == "__main__":
 
     # Setup configuration
     setup(
+        version=get_vllm_omni_version(),
         install_requires=install_requires,
     )

tests/benchmarks/test_serve_cli.py

@@ -3,36 +3,18 @@ from pathlib import Path
 
 import pytest
 
-from tests.conftest import OmniServer
 from tests.utils import hardware_test
 
-models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
-stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen3_omni_ci.yaml")]
+models = ["Qwen/Qwen2.5-Omni-7B"]
+stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen2_5_omni_ci.yaml")]
 
 # Create parameter combinations for model and stage config
 test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
 
 
-@pytest.fixture(scope="module")
-def omni_server(request):
-    """Start vLLM-Omni server as a subprocess with actual model weights.
-    Uses session scope so the server starts only once for the entire test session.
-    Multi-stage initialization can take 10-20+ minutes.
-    """
-    model, stage_config_path = request.param
-
-    print(f"Starting OmniServer with model: {model}")
-    print("This may take 10-20+ minutes for initialization...")
-
-    with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server:
-        print("OmniServer started successfully")
-        yield server
-        print("OmniServer stopped")
-
-
 @pytest.mark.core_model
 @pytest.mark.benchmark
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@hardware_test(res={"cuda": "L4"}, num_cards=3)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_bench_serve_chat(omni_server):
     command = [

tests/comfyui/conftest.py

@@ -0,0 +1,82 @@
+"""
+Conftest for ComfyUI-vLLM-Omni tests.
+
+This module sets up the test environment by:
+1. Adding the ComfyUI plugin to Python path
+2. Mocking comfy_api.input module (AudioInput, VideoInput) since comfyui is not installed
+3. Mocking comfy_extras.nodes_audio module
+"""
+
+import os
+import sys
+from typing import BinaryIO, TypedDict
+from unittest.mock import MagicMock
+
+
+def pytest_configure(config):
+    """
+    Called after command line options have been parsed and before test collection.
+    This is the right place to set up sys.path and mock modules.
+    """
+    _setup_comfyui_test_environment()
+
+
+def _setup_comfyui_test_environment():
+    """Set up the test environment for ComfyUI plugin testing."""
+    # === Add ComfyUI plugin path to allow importing comfyui_vllm_omni ===
+    _COMFYUI_PLUGIN_PATH = os.path.abspath(
+        os.path.join(os.path.dirname(__file__), "..", "..", "apps", "ComfyUI-vLLM-Omni")
+    )
+    if not os.path.isdir(_COMFYUI_PLUGIN_PATH):
+        raise FileNotFoundError(
+            f"ComfyUI plugin not found at {_COMFYUI_PLUGIN_PATH}. "
+            "If it is moved elsewhere, please update the path in this conftest.py."
+        )
+    if _COMFYUI_PLUGIN_PATH not in sys.path:
+        sys.path.insert(0, _COMFYUI_PLUGIN_PATH)
+
+    # Import torch after changing import paths. (To be used later)
+    import torch
+
+    # === Mock ComfyUI internal modules (comfy_api & comfy_extras) and "import" them to sys.module ===
+    class AudioInput(TypedDict):
+        """Mock AudioInput TypedDict from comfy_api.input"""
+
+        waveform: torch.Tensor  # Shape: (B, C, T)
+        sample_rate: int
+
+    class VideoInput:
+        """Mock VideoInput class from comfy_api.input"""
+
+        def __init__(self, data: bytes = b"mock_video_data"):
+            self._data = data
+
+        def save_to(self, file: str | BinaryIO):
+            """Save video data to file or file-like object."""
+            if isinstance(file, str):
+                print("Called VideoInput.save_to with file path. Saving to a path is no-op in tests.")
+            else:
+                file.write(self._data)
+
+    mock_comfy_api = MagicMock()
+    mock_comfy_api_input = MagicMock()
+    mock_comfy_api_input.AudioInput = AudioInput
+    mock_comfy_api_input.VideoInput = VideoInput
+    mock_comfy_api.input = mock_comfy_api_input
+
+    def mock_load(_: str | BinaryIO):
+        """Mock nodes_audio.load that returns a waveform tensor (channels, samples) and sample rate."""
+        waveform = torch.zeros((1, 24000), dtype=torch.float32)
+        sample_rate = 24000
+        return waveform, sample_rate
+
+    mock_comfy_extras = MagicMock()
+    mock_nodes_audio = MagicMock()
+    mock_nodes_audio.load = mock_load
+    mock_comfy_extras.nodes_audio = mock_nodes_audio
+
+    # Install mock modules BEFORE importing any comfyui_vllm_omni code
+    sys.modules["comfy_api"] = mock_comfy_api
+    sys.modules["comfy_api.input"] = mock_comfy_api_input
+    sys.modules["comfy_extras"] = mock_comfy_extras
+    sys.modules["comfy_extras.nodes_audio"] = mock_nodes_audio

tests/comfyui/test_comfyui_integration.py

@@ -0,0 +1,567 @@
+"""
+Integration tests for ComfyUI nodes that use the Omni API client, with a mocked AsyncOmni and a real API server running in a background process.
+These tests cover the integration between ComfyUI node and the API server, without actual model inference logic.
+It ensures that
+1. Changes made to the API (e.g., request and response formats) do not break the ComfyUI frontend that use it.
+2. The sampling parameters are correctly passed from the node to AsyncOmni through the API layer.
+"""
+
+import multiprocessing
+import time
+import traceback
+from collections.abc import Iterable, Sequence
+from enum import StrEnum, auto
+from typing import Any, NamedTuple
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+import requests
+import torch
+from comfy_api.input import AudioInput, VideoInput
+from comfyui_vllm_omni.nodes import (
+    VLLMOmniGenerateImage,
+    VLLMOmniTTS,
+    VLLMOmniUnderstanding,
+    VLLMOmniVoiceClone,
+)
+from comfyui_vllm_omni.utils.types import AutoregressionSamplingParams, DiffusionSamplingParams
+from PIL import Image
+from vllm import SamplingParams
+from vllm.outputs import CompletionOutput, RequestOutput
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.entrypoints.cli.serve import OmniServeCommand
+from vllm_omni.inputs.data import OmniSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+
+class ServerCase(NamedTuple):
+    """Parametrizing the model to serve."""
+
+    served_model: str
+    stage_list: list
+    stage_configs: list[dict]
+    outputs: list[OmniRequestOutput]
+
+
+class SamplingCase(NamedTuple):
+    """Parametrizing the input sampling parameters."""
+
+    kind: "SamplingKind"
+    sampling_params: dict | list[dict] | None
+
+
+class SamplingKind(StrEnum):
+    IMAGE_NONE = auto()
+    IMAGE_DIFFUSION_SINGLE = auto()
+    UNDERSTANDING_NONE = auto()
+    UNDERSTANDING_AR_LIST = auto()
+    TTS_NONE = auto()
+    TTS_DIFFUSION_SINGLE = auto()
+
+
+# Pre-defined arguments to be used in function calls during the tests
+IMAGE_WIDTH = 64
+IMAGE_HEIGHT = 64
+DIFFUSION_SINGLE_SAMPLING_PARAMS = DiffusionSamplingParams(
+    {
+        "n": 2,
+        "num_inference_steps": 30,
+        "guidance_scale": 6.0,
+        "true_cfg_scale": 1.5,
+    }
+)
+
+AR_LIST_SAMPLING_PARAMS = [
+    AutoregressionSamplingParams(
+        {
+            "max_tokens": 64,
+            "temperature": 0.6,
+            "top_p": 0.9,
+            "repetition_penalty": 1.0,
+            "seed": 21,
+        }
+    ),
+    AutoregressionSamplingParams(
+        {
+            "max_tokens": 96,
+            "temperature": 0.75,
+            "top_p": 0.85,
+            "repetition_penalty": 1.05,
+            "seed": 22,
+        }
+    ),
+    AutoregressionSamplingParams(
+        {
+            "max_tokens": 128,
+            "temperature": 0.8,
+            "top_p": 0.8,
+            "repetition_penalty": 1.1,
+            "seed": 23,
+        }
+    ),
+]
+
+
+def _build_image_output(size: tuple[int, int] = (IMAGE_WIDTH, IMAGE_HEIGHT), color: str = "red") -> Image.Image:
+    return Image.new("RGB", size, color=color)
+
+
+def _build_text_output(text: str = "This is a test response.") -> OmniRequestOutput:
+    completion_output = CompletionOutput(
+        index=0,
+        text=text,
+        token_ids=[1, 2, 3],
+        cumulative_logprob=0.0,
+        logprobs=None,
+        finish_reason="stop",
+        stop_reason=None,
+    )
+    request_output = RequestOutput(
+        request_id="test_req_text",
+        prompt="test prompt",
+        prompt_token_ids=[1, 2, 3],
+        prompt_logprobs=None,
+        outputs=[completion_output],
+        finished=True,
+        metrics=None,
+        lora_request=None,
+    )
+    return OmniRequestOutput(
+        request_id="test_req_text",
+        finished=True,
+        final_output_type="text",
+        request_output=request_output,
+    )
+
+
+def _build_audio_chat_output(num_samples: int = 24000) -> OmniRequestOutput:
+    completion_output = CompletionOutput(
+        index=0,
+        text="",
+        token_ids=[],
+        cumulative_logprob=0.0,
+        logprobs=None,
+        finish_reason="stop",
+        stop_reason=None,
+    )
+    completion_output.multimodal_output = {"audio": [torch.zeros(1, num_samples)]}
+    request_output = RequestOutput(
+        request_id="test_req_audio_chat",
+        prompt="test prompt",
+        prompt_token_ids=[1, 2, 3],
+        prompt_logprobs=None,
+        outputs=[completion_output],
+        finished=True,
+        metrics=None,
+        lora_request=None,
+    )
+    return OmniRequestOutput(
+        request_id="test_req_audio_chat",
+        finished=True,
+        final_output_type="audio",
+        request_output=request_output,
+    )
+
+
+def _build_audio_speech_output(num_samples: int = 24000) -> OmniRequestOutput:
+    return OmniRequestOutput.from_diffusion(
+        request_id="test_req_audio_speech",
+        images=[],
+        multimodal_output={"audio": torch.zeros(num_samples), "sr": 24000},
+        final_output_type="audio",
+    )
+
+
+def _build_diffusion_image_output_for_images_endpoint() -> OmniRequestOutput:
+    return OmniRequestOutput.from_diffusion(
+        request_id="test_req_img_dalle",
+        images=[_build_image_output()],
+        final_output_type="image",
+    )
+
+
+def _build_diffusion_image_output_for_chat_endpoint() -> OmniRequestOutput:
+    request_output = MagicMock()
+    request_output.images = [_build_image_output(color="blue")]
+    request_output.finished = True
+    return OmniRequestOutput(
+        request_id="test_req_img_chat",
+        finished=True,
+        final_output_type="image",
+        request_output=request_output,
+    )
+
+
+def _assert_sampling_param_values(received: OmniSamplingParams, expected: dict[str, Any]):
+    for key, expected_value in expected.items():
+        actual_value = getattr(received, key, None)
+        assert actual_value == expected_value, (
+            f"Expected sampling param '{key}'={expected_value}, got {actual_value}. The received sampling params: {received}"
+        )
+
+
+def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: SamplingCase, server_case: ServerCase):
+    async def _mock_generate(*args, **kwargs):
+        received_sampling_params_list: Sequence[OmniSamplingParams] | None = (
+            args[2] if len(args) > 2 else kwargs.get("sampling_params_list")
+        )
+
+        assert received_sampling_params_list is not None, (
+            "In the current codebase, the API layer always provides not-None sampling parameter list when calling `AsyncOmni.generate`"
+            "This test also uses this assumption for now."
+            "If this assertion fails, it means the API layer has changed and this test needs to be updated accordingly."
+            "It does not necessarily mean there is a bug, because `AsyncOmni.generate` does allow sampling_params_list to be None."
+        )
+        assert isinstance(received_sampling_params_list, Sequence), "sampling_params_list should be a Sequence"
+
+        if sampling_case.kind is SamplingKind.IMAGE_NONE:
+            assert len(received_sampling_params_list) == 1
+            _assert_sampling_param_values(
+                received_sampling_params_list[0],
+                {
+                    "width": IMAGE_WIDTH,
+                    "height": IMAGE_HEIGHT,
+                },
+            )
+        elif sampling_case.kind is SamplingKind.IMAGE_DIFFUSION_SINGLE:
+            assert len(received_sampling_params_list) == 1
+            expected = DIFFUSION_SINGLE_SAMPLING_PARAMS.copy()
+            expected["num_outputs_per_prompt"] = expected.pop("n")  # convert from n to num_outputs_per_prompt
+            _assert_sampling_param_values(
+                received_sampling_params_list[0],
+                {
+                    "width": IMAGE_WIDTH,
+                    "height": IMAGE_HEIGHT,
+                    **expected,
+                },
+            )
+        elif sampling_case.kind is SamplingKind.UNDERSTANDING_NONE:
+            assert len(received_sampling_params_list) == 3
+        elif sampling_case.kind is SamplingKind.UNDERSTANDING_AR_LIST:
+            assert len(received_sampling_params_list) == 3
+            for i, expected in enumerate(AR_LIST_SAMPLING_PARAMS):
+                _assert_sampling_param_values(received_sampling_params_list[i], expected)
+        elif sampling_case.kind in {SamplingKind.TTS_NONE, SamplingKind.TTS_DIFFUSION_SINGLE}:
+            assert len(received_sampling_params_list) == 1
+        else:
+            raise AssertionError(f"Unknown sampling case: {sampling_case.kind}")
+
+        for output in outputs:
+            yield output
+
+    return _mock_generate
+
+
+@pytest.fixture
+def server_case(request) -> ServerCase:
+    return request.param
+
+
+@pytest.fixture
+def sampling_case(request) -> SamplingCase:
+    return request.param
+
+
+@pytest.fixture
+def mock_async_omni(server_case: ServerCase, sampling_case: SamplingCase):
+    async def _mock_preprocess_chat(self, *args, **kwargs):
+        return ([{"role": "user", "content": "test"}], [{"prompt": "test prompt"}])
+
+    # Need to mock AsyncOmni itself (not only its generate method) because
+    # 1. The API layer uses its stage_list and stage_configs attributes
+    # 2. Its __init__ method has slow side effects (model & config loading).
+    with (
+        patch("vllm_omni.entrypoints.openai.api_server.AsyncOmni") as MockAsyncOmni,
+        patch(
+            "vllm_omni.entrypoints.openai.serving_chat.OmniOpenAIServingChat._preprocess_chat",
+            new=_mock_preprocess_chat,
+        ),
+    ):
+        mock_instance = AsyncMock()
+        mock_instance.generate = _build_mock_outputs(server_case.outputs, sampling_case, server_case)
+
+        mock_instance.stage_list = server_case.stage_list
+        mock_instance.stage_configs = server_case.stage_configs
+        mock_instance.default_sampling_params_list = [
+            SamplingParams() if stage.get("stage_type") != "diffusion" else MagicMock()
+            for stage in server_case.stage_configs
+        ]
+        mock_instance.errored = False
+        mock_instance.dead_error = RuntimeError("Mock engine error")
+        mock_instance.model_config = MagicMock(max_model_len=4096, io_processor_plugin=None)
+        mock_instance.io_processor = MagicMock()
+        mock_instance.input_processor = MagicMock()
+        mock_instance.shutdown = MagicMock()
+        mock_instance.get_vllm_config = AsyncMock(return_value=None)
+        mock_instance.get_supported_tasks = AsyncMock(return_value=["generate"])
+        mock_instance.get_tokenizer = AsyncMock(return_value=None)
+
+        MockAsyncOmni.return_value = mock_instance
+        yield MockAsyncOmni
+
+
+@pytest.fixture
+def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni):
+    """Set up a API server in background process from command line with parametrized model name and mocked AsyncOmni."""
+    parser = FlexibleArgumentParser()
+    subparsers = parser.add_subparsers(dest="command")
+    cmd = OmniServeCommand()
+    cmd.subparser_init(subparsers)
+
+    port = unused_tcp_port_factory()
+    args = parser.parse_args(["serve", server_case.served_model, "--omni", "--port", str(port)])
+
+    def run_server():
+        try:
+            cmd.cmd(args)
+        except Exception:
+            traceback.print_exc()
+
+    server_process = multiprocessing.Process(target=run_server)
+    server_process.start()
+
+    # Wait for the server to be ready by polling the health endpoint.
+    wait_time = 30
+    wait_poll_interval = 1
+    for _ in range(wait_time // wait_poll_interval):
+        try:
+            response = requests.get(f"http://127.0.0.1:{port}/health")
+            if response.status_code == 200:
+                break
+        except requests.ConnectionError:
+            time.sleep(wait_poll_interval)
+    else:
+        if server_process.is_alive():
+            server_process.terminate()
+            server_process.join(timeout=5)
+            if server_process.is_alive():
+                server_process.kill()
+                server_process.join(timeout=5)
+        pytest.fail(f"API server failed to start within {wait_time} seconds")
+
+    yield f"http://127.0.0.1:{port}/v1"
+
+    if server_process.is_alive():
+        server_process.terminate()
+    server_process.join(timeout=10)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "server_case,model,image_input",
+    [
+        pytest.param(
+            ServerCase(
+                served_model="Tongyi-MAI/Z-Image-Turbo",
+                stage_list=["diffusion"],
+                stage_configs=[{"stage_type": "diffusion"}],
+                outputs=[_build_diffusion_image_output_for_images_endpoint()],
+            ),
+            "Tongyi-MAI/Z-Image-Turbo",
+            False,
+            id="text-to-image-dalle-endpoint",
+        ),
+        pytest.param(
+            ServerCase(
+                served_model="ByteDance-Seed/BAGEL-7B-MoT",
+                stage_list=["diffusion"],
+                stage_configs=[{"stage_type": "diffusion"}],
+                outputs=[_build_diffusion_image_output_for_chat_endpoint()],
+            ),
+            "ByteDance-Seed/BAGEL-7B-MoT",
+            False,
+            id="text-to-image-bagel-chat-endpoint",
+        ),
+        pytest.param(
+            ServerCase(
+                served_model="Qwen/Qwen-Image-Edit",
+                stage_list=["diffusion"],
+                stage_configs=[{"stage_type": "diffusion"}],
+                outputs=[_build_diffusion_image_output_for_images_endpoint()],
+            ),
+            "Qwen/Qwen-Image-Edit",
+            True,
+            id="image-to-image-dalle-endpoint",
+        ),
+        pytest.param(
+            ServerCase(
+                served_model="ByteDance-Seed/BAGEL-7B-MoT",
+                stage_list=["diffusion"],
+                stage_configs=[{"stage_type": "diffusion"}],
+                outputs=[_build_diffusion_image_output_for_chat_endpoint()],
+            ),
+            "ByteDance-Seed/BAGEL-7B-MoT",
+            True,
+            id="image-to-image-bagel-chat-endpoint",
+        ),
+    ],
+    indirect=["server_case"],
+)
+@pytest.mark.parametrize(
+    "sampling_case",
+    [
+        pytest.param(SamplingCase(kind=SamplingKind.IMAGE_NONE, sampling_params=None), id="no-sampling-params"),
+        pytest.param(
+            SamplingCase(kind=SamplingKind.IMAGE_DIFFUSION_SINGLE, sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS),
+            id="single-diffusion-sampling-params",
+        ),
+    ],
+    indirect=["sampling_case"],
+)
+async def test_image_generation_node(api_server: str, model: str, image_input: bool, sampling_case: SamplingCase):
+    node = VLLMOmniGenerateImage()
+
+    kwargs = {
+        "url": api_server,
+        "model": model,
+        "prompt": "A beautiful sunset",
+        "width": IMAGE_WIDTH,
+        "height": IMAGE_HEIGHT,
+    }
+    if image_input:
+        kwargs["image"] = torch.zeros((1, IMAGE_WIDTH, IMAGE_HEIGHT, 3), dtype=torch.float32)
+    if sampling_case.sampling_params is not None:
+        kwargs["sampling_params"] = sampling_case.sampling_params
+    print(f"!!!!!! Calling {model} node.generate with kwargs: {sampling_case.sampling_params}")
+
+    result = await node.generate(**kwargs)
+
+    assert isinstance(result, tuple)
+    assert len(result) == 1
+    assert isinstance(result[0], torch.Tensor)
+    assert result[0].shape == (1, IMAGE_WIDTH, IMAGE_HEIGHT, 3)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "server_case",
+    [
+        pytest.param(
+            ServerCase(
+                served_model="Qwen/Qwen2.5-Omni-7B",
+                stage_list=[
+                    MagicMock(is_comprehension=True, model_stage="llm"),
+                    MagicMock(is_comprehension=False, model_stage="llm"),
+                    MagicMock(is_comprehension=False, model_stage="llm"),
+                ],
+                stage_configs=[
+                    {"stage_type": "llm"},
+                    {"stage_type": "llm"},
+                    {"stage_type": "llm"},
+                ],
+                outputs=[_build_audio_chat_output(), _build_text_output("Understanding response")],
+            ),
+            id="multimodal-understanding",
+        )
+    ],
+    indirect=["server_case"],
+)
+@pytest.mark.parametrize(
+    "sampling_case",
+    [
+        pytest.param(SamplingCase(kind=SamplingKind.UNDERSTANDING_NONE, sampling_params=None), id="no-sampling-params"),
+        pytest.param(
+            SamplingCase(kind=SamplingKind.UNDERSTANDING_AR_LIST, sampling_params=AR_LIST_SAMPLING_PARAMS),
+            id="ar-sampling-params-list",
+        ),
+    ],
+    indirect=["sampling_case"],
+)
+async def test_understanding_node(api_server: str, sampling_case: SamplingCase):
+    node = VLLMOmniUnderstanding()
+
+    image = torch.zeros((1, IMAGE_WIDTH, IMAGE_HEIGHT, 3), dtype=torch.float32)
+    video = VideoInput(b"mock_video_for_test")  # type: ignore[reportAbstractUsage]
+    audio: AudioInput = {"waveform": torch.zeros((1, 1, 24000), dtype=torch.float32), "sample_rate": 24000}
+
+    text_response, audio_response = await node.generate(
+        url=api_server,
+        model="Qwen/Qwen2.5-Omni-7B",
+        prompt="Describe all modalities.",
+        image=image,
+        audio=audio,
+        video=video,
+        sampling_params=sampling_case.sampling_params,
+        output_text=True,
+        output_audio=True,
+        use_audio_in_video=True,
+    )
+
+    assert text_response == "Understanding response"
+    assert isinstance(audio_response, dict)
+    assert audio_response["sample_rate"] == 24000
+    assert isinstance(audio_response["waveform"], torch.Tensor)
+    assert audio_response["waveform"].shape == (1, 1, 24000)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "server_case,node_cls,call_kwargs",
+    [
+        pytest.param(
+            ServerCase(
+                served_model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+                stage_list=["llm"],
+                stage_configs=[{"stage_type": "llm"}],
+                outputs=[_build_audio_speech_output()],
+            ),
+            VLLMOmniTTS,
+            {
+                "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+                "input": "Hello from TTS test",
+                "voice": "Vivian",
+                "response_format": "wav",
+                "speed": 1.0,
+                "model_specific_params": None,
+            },
+            id="tts",
+        ),
+        pytest.param(
+            ServerCase(
+                served_model="Qwen/Qwen3-TTS-12Hz-1.7B-Base",
+                stage_list=["llm"],
+                stage_configs=[{"stage_type": "llm"}],
+                outputs=[_build_audio_speech_output()],
+            ),
+            VLLMOmniVoiceClone,
+            {
+                "model": "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
+                "input": "Hello from voice clone test",
+                "voice": "Vivian",
+                "response_format": "wav",
+                "speed": 1.0,
+                "ref_audio": {"waveform": torch.zeros((1, 1, 24000), dtype=torch.float32), "sample_rate": 24000},
+                "ref_text": "Reference transcript",
+                "x_vector_only_mode": False,
+                "model_specific_params": None,
+            },
+            id="tts-voice-clone",
+        ),
+    ],
+    indirect=["server_case"],
+)
+@pytest.mark.parametrize(
+    "sampling_case",
+    [
+        pytest.param(SamplingCase(kind=SamplingKind.TTS_NONE, sampling_params=None), id="no-sampling-params"),
+        pytest.param(
+            SamplingCase(kind=SamplingKind.TTS_DIFFUSION_SINGLE, sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS),
+            id="single-diffusion-sampling-params",
+        ),
+    ],
+    indirect=["sampling_case"],
+)
+async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_case: SamplingCase):
+    node = node_cls()
+    actual_kwargs = dict(call_kwargs)
+    if sampling_case.sampling_params is not None:
+        actual_kwargs["model_specific_params"] = sampling_case.sampling_params
+    result = await node.generate(url=api_server, **actual_kwargs)
+
+    assert isinstance(result, tuple)
+    assert len(result) == 1
+    assert isinstance(result[0], dict)
+    assert result[0]["sample_rate"] == 24000
+    assert isinstance(result[0]["waveform"], torch.Tensor)
+    assert result[0]["waveform"].shape == (1, 1, 24000)

tests/diffusion/attention/test_flash_attn.py

@@ -15,6 +15,9 @@ import torch
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
 from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
+from vllm_omni.platforms import current_omni_platform
+
+is_gpu = current_omni_platform.is_cuda_alike() or current_omni_platform.is_xpu()
 
 
 def create_attention_mask(batch_size: int, seq_len: int, valid_len: int, device: torch.device) -> torch.Tensor:
@@ -56,7 +59,7 @@ def pad_tensor(tensor: torch.Tensor, target_seq_len: int, pad_value: float = 0.0
     return torch.cat([tensor, padding], dim=1)
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="FlashAttention requires CUDA")
+@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU")
 def test_padding_equivalence():
     """
     Case 1: Test that padded and unpadded inputs produce similar outputs.
@@ -68,7 +71,7 @@ def test_padding_equivalence():
 
     Expected: Output A and Output B should be very close.
     """
-    device = torch.device("cuda")
+    device = torch.device(current_omni_platform.device_type)
     dtype = torch.bfloat16
 
     # Configuration
@@ -149,7 +152,7 @@ def test_padding_equivalence():
     print("✓ Case 1 PASSED: Padded and unpadded outputs are very close!")
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="FlashAttention requires CUDA")
+@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU")
 def test_fa_vs_sdpa():
     """
     Case 2: Compare FlashAttention and SDPA backends with padding.
@@ -162,7 +165,7 @@ def test_fa_vs_sdpa():
 
     Expected: FA and SDPA outputs should be very close.
     """
-    device = torch.device("cuda")
+    device = torch.device(current_omni_platform.device_type)
     dtype = torch.bfloat16
 
     # Configuration
@@ -265,8 +268,8 @@ if __name__ == "__main__":
     print("Running FlashAttention Padding Tests...")
     print("=" * 60)
 
-    # Try to run CUDA tests
-    if torch.cuda.is_available():
+    # Try to run tests
+    if is_gpu:
         try:
             print("\n[Running Case 1: Padding Equivalence for FA]")
             test_padding_equivalence()

tests/diffusion/distributed/test_hsdp.py

@@ -0,0 +1,225 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities.
+
+These tests verify HSDP configuration logic without requiring a distributed environment.
+"""
+
+import pytest
+import torch.nn as nn
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu]
+
+
+class TestHSDPInferenceConfig:
+    """Tests for HSDPInferenceConfig dataclass."""
+
+    def test_default_values(self):
+        """Test default configuration values."""
+        config = HSDPInferenceConfig()
+        assert config.enabled is False
+        assert config.hsdp_replicate_size == 1
+        assert config.hsdp_shard_size == -1
+
+    def test_custom_values(self):
+        """Test custom configuration values."""
+        config = HSDPInferenceConfig(
+            enabled=True,
+            hsdp_replicate_size=2,
+            hsdp_shard_size=4,
+        )
+        assert config.enabled is True
+        assert config.hsdp_replicate_size == 2
+        assert config.hsdp_shard_size == 4
+
+
+class TestDiffusionParallelConfigHSDP:
+    """Tests for HSDP settings in DiffusionParallelConfig."""
+
+    def test_hsdp_disabled_by_default(self):
+        """HSDP should be disabled by default."""
+        config = DiffusionParallelConfig()
+        assert config.use_hsdp is False
+        assert config.hsdp_shard_size == -1
+        assert config.hsdp_replicate_size == 1
+
+    def test_hsdp_auto_shard_size(self):
+        """Test auto-calculation of hsdp_shard_size when use_hsdp=True."""
+        # ulysses_degree=4 -> world_size=4
+        # hsdp_shard_size should be auto-calculated as 4 // 1 = 4
+        config = DiffusionParallelConfig(
+            ulysses_degree=4,
+            use_hsdp=True,
+        )
+        assert config.world_size == 4
+        assert config.hsdp_shard_size == 4
+        assert config.hsdp_replicate_size == 1
+
+    def test_hsdp_auto_shard_size_fails_standalone(self):
+        """Test that auto-calculate fails when other parallelism is all 1."""
+        # When all other parallelism is 1, cannot auto-calculate
+        # User must specify hsdp_shard_size explicitly
+        with pytest.raises(ValueError, match="Cannot auto-calculate hsdp_shard_size"):
+            DiffusionParallelConfig(
+                use_hsdp=True,
+                # All other parallelism defaults to 1
+            )
+
+    def test_hsdp_standalone_mode(self):
+        """Test standalone HSDP (HSDP without other parallelism)."""
+        # Standalone HSDP: all other parallelism=1, explicit shard_size
+        config = DiffusionParallelConfig(
+            use_hsdp=True,
+            hsdp_shard_size=4,  # Explicit shard size
+            hsdp_replicate_size=1,
+        )
+        # world_size should be determined by HSDP
+        assert config.world_size == 4
+        assert config.hsdp_shard_size == 4
+        assert config.hsdp_replicate_size == 1
+
+    def test_hsdp_standalone_with_replicate(self):
+        """Test standalone HSDP with replication."""
+        config = DiffusionParallelConfig(
+            use_hsdp=True,
+            hsdp_shard_size=4,
+            hsdp_replicate_size=2,
+        )
+        # world_size = shard_size * replicate_size
+        assert config.world_size == 8
+        assert config.hsdp_shard_size == 4
+        assert config.hsdp_replicate_size == 2
+
+    def test_hsdp_with_replicate(self):
+        """Test HSDP with replication (hybrid mode) combined with other parallelism."""
+        # world_size=8, replicate=2 -> shard_size should be 4
+        config = DiffusionParallelConfig(
+            ulysses_degree=8,
+            use_hsdp=True,
+            hsdp_replicate_size=2,
+        )
+        assert config.world_size == 8
+        assert config.hsdp_shard_size == 4
+        assert config.hsdp_replicate_size == 2
+
+    def test_hsdp_explicit_shard_size_valid(self):
+        """Test explicit hsdp_shard_size that matches world_size."""
+        config = DiffusionParallelConfig(
+            ulysses_degree=4,
+            use_hsdp=True,
+            hsdp_shard_size=4,
+            hsdp_replicate_size=1,
+        )
+        assert config.hsdp_shard_size == 4
+
+    def test_hsdp_explicit_shard_size_invalid(self):
+        """Test that invalid HSDP dimensions raise an error when combined with other parallelism."""
+        with pytest.raises(ValueError, match="HSDP dimensions"):
+            DiffusionParallelConfig(
+                ulysses_degree=4,  # world_size=4
+                use_hsdp=True,
+                hsdp_shard_size=3,  # 1 * 3 != 4
+                hsdp_replicate_size=1,
+            )
+
+    def test_hsdp_replicate_size_exceeds_world_size(self):
+        """Test that replicate_size > world_size raises an error."""
+        with pytest.raises(ValueError, match="replicate_size.*must evenly divide world_size"):
+            DiffusionParallelConfig(
+                ulysses_degree=4,  # world_size=4
+                use_hsdp=True,
+                hsdp_replicate_size=8,  # 8 > 4, invalid
+            )
+
+    def test_hsdp_combined_world_size(self):
+        """Test that combined HSDP matches other parallelism world_size."""
+        config_no_hsdp = DiffusionParallelConfig(ulysses_degree=4)
+        config_with_hsdp = DiffusionParallelConfig(
+            ulysses_degree=4,
+            use_hsdp=True,
+            hsdp_shard_size=4,
+        )
+        # When combined with other parallelism, world_size should match
+        assert config_no_hsdp.world_size == config_with_hsdp.world_size == 4
+
+    def test_hsdp_standalone_world_size(self):
+        """Test that standalone HSDP determines world_size."""
+        config_hsdp = DiffusionParallelConfig(
+            use_hsdp=True,
+            hsdp_shard_size=8,
+        )
+        # Standalone HSDP: world_size is determined by HSDP
+        assert config_hsdp.world_size == 8
+
+    def test_hsdp_cannot_use_with_tp(self):
+        """Test that HSDP and Tensor Parallelism cannot be used together."""
+        with pytest.raises(ValueError, match="cannot be used with Tensor Parallelism"):
+            DiffusionParallelConfig(
+                tensor_parallel_size=2,
+                use_hsdp=True,
+                hsdp_shard_size=4,
+            )
+
+    def test_from_dict_with_hsdp(self):
+        """Test creating config from dict with HSDP settings."""
+        config = DiffusionParallelConfig.from_dict(
+            {
+                "ulysses_degree": 4,
+                "use_hsdp": True,
+                "hsdp_replicate_size": 2,
+            }
+        )
+        assert config.use_hsdp is True
+        assert config.hsdp_replicate_size == 2
+        assert config.hsdp_shard_size == 2  # auto: 4 // 2
+
+
+class TestHSDPShardConditions:
+    """Tests for _hsdp_shard_conditions matching logic."""
+
+    @staticmethod
+    def _is_transformer_block(name: str, module: nn.Module) -> bool:
+        """Example shard condition matching transformer blocks."""
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    def test_condition_matches_blocks(self):
+        """Test that condition matches transformer block patterns."""
+        cond = self._is_transformer_block
+        # Should match
+        assert cond("blocks.0", nn.Linear(10, 10)) is True
+        assert cond("blocks.15", nn.Linear(10, 10)) is True
+        assert cond("transformer.blocks.0", nn.Linear(10, 10)) is True
+        # Should not match
+        assert cond("blocks", nn.Linear(10, 10)) is False
+        assert cond("blocks.norm", nn.Linear(10, 10)) is False
+        assert cond("embeddings", nn.Linear(10, 10)) is False
+
+    def test_model_with_shard_conditions(self):
+        """Test model with _hsdp_shard_conditions attribute."""
+
+        class MockModel(nn.Module):
+            @staticmethod
+            def _is_block(name: str, module: nn.Module) -> bool:
+                return name.startswith("blocks.") and name.split(".")[-1].isdigit()
+
+            _hsdp_shard_conditions = [_is_block]
+
+            def __init__(self):
+                super().__init__()
+                self.blocks = nn.ModuleList([nn.Linear(10, 10) for _ in range(2)])
+
+        model = MockModel()
+        conditions = getattr(model, "_hsdp_shard_conditions", None)
+        assert conditions is not None
+        assert len(conditions) == 1
+
+        # Verify conditions work on actual model modules
+        matched = []
+        for name, module in model.named_modules():
+            if any(cond(name, module) for cond in conditions):
+                matched.append(name)
+        assert "blocks.0" in matched
+        assert "blocks.1" in matched

tests/diffusion/test_diffusers_loader.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _DummyPipelineModel(nn.Module):
+    def __init__(self, *, source_prefix: str):
+        super().__init__()
+        self.transformer = nn.Linear(2, 2, bias=False)
+        self.vae = nn.Linear(2, 2, bias=False)
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path="dummy",
+                subfolder="transformer",
+                revision=None,
+                prefix=source_prefix,
+                fall_back_to_pt=True,
+            )
+        ]
+
+    def load_weights(self, weights):
+        params = dict(self.named_parameters())
+        loaded: set[str] = set()
+        for name, tensor in weights:
+            if name not in params:
+                continue
+            params[name].data.copy_(tensor.to(dtype=params[name].dtype))
+            loaded.add(name)
+        return loaded
+
+
+def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoader:
+    loader = object.__new__(DiffusersPipelineLoader)
+    loader.counter_before_loading_weights = 0.0
+    loader.counter_after_loading_weights = 0.0
+
+    def _iter_weights(_model):
+        for name in weight_names:
+            yield name, torch.zeros((2, 2))
+
+    loader.get_all_weights = _iter_weights  # type: ignore[assignment]
+    return loader
+
+
+def test_strict_check_only_validates_source_prefix_parameters():
+    model = _DummyPipelineModel(source_prefix="transformer.")
+    loader = _make_loader_with_weights(["transformer.weight"])
+
+    # Should not require VAE parameters because they are outside weights_sources.
+    loader.load_weights(model)
+
+
+def test_strict_check_raises_when_source_parameters_are_missing():
+    model = _DummyPipelineModel(source_prefix="transformer.")
+    loader = _make_loader_with_weights([])
+
+    with pytest.raises(ValueError, match="transformer.weight"):
+        loader.load_weights(model)
+
+
+def test_empty_source_prefix_keeps_full_model_strict_check():
+    model = _DummyPipelineModel(source_prefix="")
+    loader = _make_loader_with_weights(["transformer.weight"])
+
+    with pytest.raises(ValueError, match="vae.weight"):
+        loader.load_weights(model)

tests/diffusion/test_diffusion_model_runner.py

@@ -54,6 +54,7 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
     runner.od_config = SimpleNamespace(
         cache_backend=cache_backend_name,
         enable_cache_dit_summary=enable_cache_dit_summary,
+        parallel_config=SimpleNamespace(use_hsdp=False),
     )
     runner.kv_transfer_manager = SimpleNamespace(receive_kv_cache=lambda req, target_device: None)
     return runner

tests/diffusion/test_multiproc_executor_concurrency.py

@@ -0,0 +1,442 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import queue
+import threading
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
+from vllm_omni.diffusion.scheduler import Scheduler
+
+pytestmark = [pytest.mark.diffusion]
+
+
+# ───────────────────────────────────────────── helpers ─────────────────────
+
+
+def _tagged_output(tag: str) -> DiffusionOutput:
+    """Return a ``DiffusionOutput`` identifiable by its *error* field."""
+    return DiffusionOutput(output=torch.tensor([0]), error=tag)
+
+
+def _mock_request(tag: str) -> Mock:
+    """Return a mock ``OmniDiffusionRequest`` identifiable by *tag*."""
+    req = Mock()
+    req.request_ids = [tag]
+    return req
+
+
+def _make_scheduler():
+    """Create a ``Scheduler`` whose *mq* / *result_mq* are backed by
+    plain ``queue.Queue`` objects (thread-safe, no real IPC).
+
+    Returns ``(scheduler, request_queue, result_queue)``.
+    """
+    sched = Scheduler()
+    sched.num_workers = 1
+    sched._lock = threading.Lock()
+
+    req_q: queue.Queue = queue.Queue()
+    res_q: queue.Queue = queue.Queue()
+
+    mock_mq = Mock()
+    mock_mq.enqueue = req_q.put
+
+    mock_rmq = Mock()
+    mock_rmq.dequeue = lambda timeout=None: res_q.get(timeout=timeout if timeout else 10)
+
+    sched.mq = mock_mq
+    sched.result_mq = mock_rmq
+    return sched, req_q, res_q
+
+
+def _make_executor(scheduler):
+    """Create a ``MultiprocDiffusionExecutor`` wired to *scheduler*
+    without launching real worker processes.
+    """
+    od_cfg = Mock()
+    od_cfg.num_gpus = 1
+    with patch.object(MultiprocDiffusionExecutor, "_init_executor"):
+        executor = MultiprocDiffusionExecutor(od_cfg)
+    executor.scheduler = scheduler
+    executor._closed = False
+    executor._processes = []
+    return executor
+
+
+def _start_worker(req_q, res_q, count=2):
+    """Simulate workers: read *count* requests from *req_q* and put
+    tagged ``DiffusionOutput``s on *res_q* (FIFO order).
+    """
+
+    def _run():
+        for _ in range(count):
+            req = req_q.get(timeout=10)
+            method = req.get("method", "")
+            args = req.get("args", ())
+            if method == "generate" and args and hasattr(args[0], "request_ids"):
+                tag = f"result_for_{args[0].request_ids[0]}"
+            elif args:
+                tag = f"result_for_{args[0]}"
+            else:
+                tag = f"result_for_{method}"
+            res_q.put(_tagged_output(tag))
+
+    t = threading.Thread(target=_run, daemon=True)
+    t.start()
+    return t
+
+
+def _inject_interleave(scheduler):
+    """Monkey-patch ``scheduler.mq.enqueue`` so that:
+
+    * The thread named **thread_a** *blocks* after its enqueue until the
+      thread named **thread_b** has finished entirely.
+    * All other threads pass through unblocked.
+
+    Returns ``(a_enqueued: Event, b_complete: Event)`` for wiring.
+    """
+    a_enqueued = threading.Event()
+    b_complete = threading.Event()
+    orig_enqueue = scheduler.mq.enqueue  # points to req_q.put
+
+    def _controlled(item):
+        orig_enqueue(item)
+        if threading.current_thread().name == "thread_a":
+            a_enqueued.set()  # tell B: "A has enqueued"
+            b_complete.wait(5)  # block A until B finishes
+
+    scheduler.mq.enqueue = _controlled
+    return a_enqueued, b_complete
+
+
+# ──────────────────── bug-reproduction: concurrent add_req ────────────────
+
+
+class TestConcurrentAddReqBug:
+    """Two concurrent ``Scheduler.add_req()`` calls swap results."""
+
+    def test_results_are_correctly_routed(self):
+        sched, req_q, res_q = _make_scheduler()
+        a_enqueued, b_complete = _inject_interleave(sched)
+        wt = _start_worker(req_q, res_q, count=2)
+
+        results: dict[str, DiffusionOutput] = {}
+
+        def _a():
+            results["A"] = sched.add_req(_mock_request("A"))
+
+        def _b():
+            a_enqueued.wait(5)  # wait for A to enqueue
+            results["B"] = sched.add_req(_mock_request("B"))
+            b_complete.set()  # release A
+
+        ta = threading.Thread(target=_a, name="thread_a")
+        tb = threading.Thread(target=_b, name="thread_b")
+        ta.start()
+        tb.start()
+        ta.join(10)
+        tb.join(10)
+        wt.join(5)
+
+        # With correct (locked) implementation both assertions hold.
+        # The bug causes them to be swapped.
+        assert results["A"].error == "result_for_A"
+        assert results["B"].error == "result_for_B"
+
+
+# ──────────────── bug-reproduction: concurrent collective_rpc ─────────────
+
+
+class TestConcurrentCollectiveRpcBug:
+    """Two concurrent ``collective_rpc()`` calls swap results."""
+
+    def test_results_are_correctly_routed(self):
+        sched, req_q, res_q = _make_scheduler()
+        executor = _make_executor(sched)
+        a_enqueued, b_complete = _inject_interleave(sched)
+        wt = _start_worker(req_q, res_q, count=2)
+
+        results: dict[str, object] = {}
+
+        def _a():
+            results["A"] = executor.collective_rpc(
+                "ping",
+                args=("call_A",),
+                unique_reply_rank=0,
+            )
+
+        def _b():
+            a_enqueued.wait(5)
+            results["B"] = executor.collective_rpc(
+                "ping",
+                args=("call_B",),
+                unique_reply_rank=0,
+            )
+            b_complete.set()
+
+        ta = threading.Thread(target=_a, name="thread_a")
+        tb = threading.Thread(target=_b, name="thread_b")
+        ta.start()
+        tb.start()
+        ta.join(10)
+        tb.join(10)
+        wt.join(5)
+
+        assert results["A"].error == "result_for_call_A"
+        assert results["B"].error == "result_for_call_B"
+
+
+# ──────── bug-reproduction: add_req vs collective_rpc concurrently ────────
+
+
+class TestConcurrentAddReqVsCollectiveRpcBug:
+    """``add_req`` and ``collective_rpc`` running concurrently swap results."""
+
+    def test_results_are_correctly_routed(self):
+        sched, req_q, res_q = _make_scheduler()
+        executor = _make_executor(sched)
+        a_enqueued, b_complete = _inject_interleave(sched)
+        wt = _start_worker(req_q, res_q, count=2)
+
+        results: dict[str, object] = {}
+
+        def _a():  # add_req path
+            results["A"] = sched.add_req(_mock_request("A"))
+
+        def _b():  # collective_rpc path
+            a_enqueued.wait(5)
+            results["B"] = executor.collective_rpc(
+                "ping",
+                args=("call_B",),
+                unique_reply_rank=0,
+            )
+            b_complete.set()
+
+        ta = threading.Thread(target=_a, name="thread_a")
+        tb = threading.Thread(target=_b, name="thread_b")
+        ta.start()
+        tb.start()
+        ta.join(10)
+        tb.join(10)
+        wt.join(5)
+
+        assert isinstance(results["A"], DiffusionOutput)
+        assert results["A"].error == "result_for_A"
+        assert results["B"].error == "result_for_call_B"
+
+
+# ─────────────── backward-compatibility (serial) tests ────────────────────
+
+
+class TestSerialOperations:
+    """Verify correct behaviour for single-threaded (serial) usage.
+
+    These tests must pass both **before** and **after** any concurrency fix
+    is applied – they guard against regressions in the basic request path.
+    """
+
+    def test_serial_add_req_returns_correct_result(self):
+        sched, req_q, res_q = _make_scheduler()
+        wt = _start_worker(req_q, res_q, count=1)
+
+        result = sched.add_req(_mock_request("X"))
+        wt.join(5)
+
+        assert isinstance(result, DiffusionOutput)
+        assert result.error == "result_for_X"
+
+    def test_serial_add_req_multiple_sequential(self):
+        sched, req_q, res_q = _make_scheduler()
+        wt = _start_worker(req_q, res_q, count=3)
+
+        for tag in ("one", "two", "three"):
+            out = sched.add_req(_mock_request(tag))
+            assert out.error == f"result_for_{tag}"
+
+        wt.join(5)
+
+    def test_serial_collective_rpc_single_rank(self):
+        sched, req_q, res_q = _make_scheduler()
+        executor = _make_executor(sched)
+        wt = _start_worker(req_q, res_q, count=1)
+
+        result = executor.collective_rpc(
+            "ping",
+            args=("Y",),
+            unique_reply_rank=0,
+        )
+        wt.join(5)
+
+        assert result.error == "result_for_Y"
+
+    def test_serial_collective_rpc_all_ranks(self):
+        """``collective_rpc`` without *unique_reply_rank* collects
+        ``num_gpus`` responses.
+        """
+        sched, req_q, res_q = _make_scheduler()
+        executor = _make_executor(sched)
+        executor.od_config.num_gpus = 2
+
+        # Pre-populate two results (simulating two workers replying)
+        res_q.put(_tagged_output("rank0"))
+        res_q.put(_tagged_output("rank1"))
+
+        results = executor.collective_rpc("ping", args=("multi",))
+
+        assert len(results) == 2
+        assert results[0].error == "rank0"
+        assert results[1].error == "rank1"
+
+    def test_serial_add_req_then_collective_rpc(self):
+        sched, req_q, res_q = _make_scheduler()
+        executor = _make_executor(sched)
+        wt = _start_worker(req_q, res_q, count=2)
+
+        gen_out = sched.add_req(_mock_request("gen"))
+        rpc_out = executor.collective_rpc(
+            "ping",
+            args=("rpc",),
+            unique_reply_rank=0,
+        )
+        wt.join(5)
+
+        assert gen_out.error == "result_for_gen"
+        assert rpc_out.error == "result_for_rpc"
+
+    def test_serial_add_req_error_propagation(self):
+        """``add_req`` should raise when the worker reports an error."""
+        sched, _, res_q = _make_scheduler()
+        # Put an error response directly
+        res_q.put({"status": "error", "error": "boom"})
+
+        with pytest.raises(RuntimeError, match="worker error"):
+            sched.add_req(_mock_request("fail"))
+
+    def test_serial_collective_rpc_error_propagation(self):
+        """``collective_rpc`` should raise when the worker reports an error."""
+        sched, _, res_q = _make_scheduler()
+        executor = _make_executor(sched)
+        res_q.put({"status": "error", "error": "kaboom"})
+
+        with pytest.raises(RuntimeError, match="kaboom"):
+            executor.collective_rpc("bad", unique_reply_rank=0)
+
+    def test_collective_rpc_closed_executor_raises(self):
+        sched, _, _ = _make_scheduler()
+        executor = _make_executor(sched)
+        executor._closed = True
+
+        with pytest.raises(RuntimeError, match="closed"):
+            executor.collective_rpc("anything")
+
+
+# ─────────── timeout regression: RPC must not block on a stalled lock ─────
+
+
+class TestCollectiveRpcTimeoutWhileLockHeld:
+    """``collective_rpc(timeout=...)`` must honour its timeout even when
+    another thread holds ``scheduler._lock`` indefinitely (e.g. a stalled
+    ``add_req`` waiting on an unresponsive worker).
+    """
+
+    def test_rpc_times_out_when_lock_held_directly(self):
+        """Simplest case: lock is manually held by another thread."""
+        sched, req_q, res_q = _make_scheduler()
+        executor = _make_executor(sched)
+
+        stall_started = threading.Event()
+
+        def _hold_lock():
+            sched._lock.acquire()
+            stall_started.set()
+            # Hold the lock far longer than the RPC timeout.
+            threading.Event().wait(30)
+            sched._lock.release()
+
+        stall_thread = threading.Thread(target=_hold_lock, daemon=True)
+        stall_thread.start()
+        stall_started.wait(5)
+
+        # collective_rpc should raise TimeoutError, NOT block forever.
+        with pytest.raises(TimeoutError):
+            executor.collective_rpc("health", timeout=0.5)
+
+    def test_rpc_times_out_when_add_req_stalled_on_worker(self):
+        """Real-world scenario the bot flagged:
+
+        ``add_req`` holds ``_lock`` while blocked on ``result_mq.dequeue()``
+        because the worker never replies.  A concurrent
+        ``collective_rpc(timeout=...)`` must still time out instead of
+        hanging forever waiting for the lock.
+        """
+        sched, req_q, res_q = _make_scheduler()
+        executor = _make_executor(sched)
+
+        add_req_blocked = threading.Event()
+
+        # Patch dequeue: signal once entered, then block indefinitely
+        # (simulates a worker that never sends a result).
+        orig_dequeue = sched.result_mq.dequeue
+
+        def _hanging_dequeue(timeout=None):
+            add_req_blocked.set()
+            # Block forever — the worker is "hung".
+            threading.Event().wait(30)
+            return orig_dequeue(timeout=timeout)
+
+        sched.result_mq.dequeue = _hanging_dequeue
+
+        # Thread running add_req — acquires the lock, enqueues, then
+        # blocks on dequeue forever (worker hang).
+        def _stalled_add_req():
+            try:
+                sched.add_req(_mock_request("stalled"))
+            except Exception:
+                pass
+
+        t = threading.Thread(target=_stalled_add_req, daemon=True)
+        t.start()
+
+        # Wait until add_req is truly inside the lock and blocking.
+        add_req_blocked.wait(5)
+
+        # collective_rpc should time out at lock acquisition, not hang.
+        with pytest.raises(TimeoutError):
+            executor.collective_rpc("health_check", timeout=0.5)
+
+    def test_rpc_without_timeout_still_waits_for_lock(self):
+        """When no timeout is given, ``collective_rpc`` should still wait
+        for the lock (blocking) — existing behaviour preserved.
+        """
+        sched, req_q, res_q = _make_scheduler()
+        executor = _make_executor(sched)
+
+        lock_released = threading.Event()
+
+        def _hold_and_release():
+            sched._lock.acquire()
+            # Hold for a short time then release.
+            threading.Event().wait(0.3)
+            sched._lock.release()
+            lock_released.set()
+
+        # Pre-populate a result so collective_rpc succeeds after lock.
+        res_q.put(_tagged_output("ok"))
+
+        t = threading.Thread(target=_hold_and_release, daemon=True)
+        t.start()
+
+        # No timeout → should block until lock is released, then succeed.
+        result = executor.collective_rpc(
+            "ping",
+            args=("wait",),
+            unique_reply_rank=0,
+        )
+        t.join(5)
+
+        assert result.error == "ok"

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -35,6 +35,7 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None)
         prompt_token_ids=[],
         num_computed_tokens=0,
         additional_information=None,
+        is_finished=lambda: status == RequestStatus.FINISHED_STOPPED,
     )
 
 
@@ -48,12 +49,11 @@ def build_adapter(monkeypatch, mocker: MockerFixture):
 
         def _fake_base_init(self, config):
             self.config = config
-            self._pending_load_reqs = {}
+            self._pending_load_reqs = deque()
             self._finished_load_reqs = set()
-            self._pending_save_reqs = {}
+            self._pending_save_reqs = deque()
             self._finished_save_reqs = set()
             self.stop_event = threading.Event()
-            self.lock = threading.Lock()
 
         monkeypatch.setattr(OmniTransferAdapterBase, "__init__", _fake_base_init)
         monkeypatch.setattr(
@@ -108,9 +108,8 @@ def test_load_poll(build_adapter):
     adapter.load_async(request)
     payload = {"code_predictor_codes": [[1]], "hidden_states": torch.tensor([[2.0]]), "finished": True}
     connector.get.return_value = (payload, 16)
-    adapter._poll_single_request("req-1")
+    adapter._poll_single_request(request)
 
-    connector.get.assert_called_once_with("1", "2", "external-1_1_0")
     assert request.additional_information == payload
     assert adapter.get_req_chunk["req-1"] == 1
     assert "req-1" in adapter._finished_load_reqs
@@ -120,17 +119,15 @@ def test_load_poll(build_adapter):
 
 def test_save_async(build_adapter):
     adapter, _ = build_adapter(stage_id=1)
-    request = SimpleNamespace(external_req_id="external-1")
+    request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")
 
     adapter.custom_process_next_stage_input_func = lambda **kwargs: {"x": [1], "finished": False}
     adapter.save_async(pooling_output=None, request=request)
     adapter.custom_process_next_stage_input_func = lambda **kwargs: {}
     adapter.save_async(pooling_output=None, request=request)
 
-    assert adapter.put_req_chunk["external-1"] == 1
-    queued = adapter._pending_save_reqs["external-1"]
-    assert len(queued) == 1
-    assert queued[0]["put_key"] == "external-1_1_0"
+    task = adapter._pending_save_reqs.popleft()
+    assert task["is_finished"] is False
 
 
 def test_update_request_payload(build_adapter):
@@ -180,3 +177,232 @@ def test_postprocess_scheduler_output(build_adapter):
     assert cached_info["cached-ready"] == {"k": "v"}
     assert cached_info["missing"] is None
     assert adapter.requests_with_ready_chunks == {"leftover"}
+
+
+# ---------------------------------------------------------------
+# Cleanup tests
+# ---------------------------------------------------------------
+
+
+def _populate_adapter_state(adapter, req_id="req-1", ext_id="ext-1"):
+    """Fill every per-request structure so cleanup can be verified."""
+    adapter.finished_requests.add(req_id)
+    adapter.get_req_chunk[req_id] = 3
+    adapter.requests_with_ready_chunks.add(req_id)
+    adapter.request_ids_mapping[req_id] = ext_id
+    adapter._pending_load_reqs.append(SimpleNamespace(request_id=req_id))
+    adapter._finished_load_reqs.add(req_id)
+
+    adapter.put_req_chunk[ext_id] = 5
+    adapter.request_payload[ext_id] = {"hidden": [1, 2]}
+    adapter.code_prompt_token_ids[ext_id] = [[10, 20]]
+
+
+def test_cleanup_clears_all_state(build_adapter):
+    """After cleanup, no per-request key should remain in any dict/set."""
+    adapter, _ = build_adapter(stage_id=1)
+    req_id, ext_id = "req-1", "ext-1"
+    _populate_adapter_state(adapter, req_id, ext_id)
+
+    adapter.cleanup(req_id, ext_id)
+
+    assert req_id not in adapter.finished_requests
+    assert req_id not in adapter.get_req_chunk
+    assert req_id not in adapter.requests_with_ready_chunks
+    assert req_id not in adapter.request_ids_mapping
+    assert all(getattr(r, "request_id", None) != req_id for r in adapter._pending_load_reqs)
+    assert req_id not in adapter._finished_load_reqs
+
+    assert ext_id not in adapter.put_req_chunk
+    assert ext_id not in adapter.request_payload
+    assert ext_id not in adapter.code_prompt_token_ids
+
+
+def test_cleanup_infers_external_id(build_adapter):
+    """When external_req_id is None, cleanup should look it up from the mapping."""
+    adapter, _ = build_adapter(stage_id=1)
+    req_id, ext_id = "req-2", "ext-2"
+    _populate_adapter_state(adapter, req_id, ext_id)
+
+    adapter.cleanup(req_id)
+
+    assert ext_id not in adapter.put_req_chunk
+    assert ext_id not in adapter.request_payload
+
+
+def test_cleanup_idempotent(build_adapter):
+    """Calling cleanup multiple times for the same (or nonexistent) request must not raise."""
+    adapter, _ = build_adapter(stage_id=1)
+
+    try:
+        adapter.cleanup("nonexistent")
+        adapter.cleanup("nonexistent")
+    except Exception as e:
+        pytest.fail(f"cleanup should be idempotent: {e}")
+
+    req_id, ext_id = "req-3", "ext-3"
+    _populate_adapter_state(adapter, req_id, ext_id)
+    adapter.cleanup(req_id, ext_id)
+
+    try:
+        adapter.cleanup(req_id, ext_id)
+    except Exception as e:
+        pytest.fail(f"second cleanup should be idempotent: {e}")
+
+
+def test_cleanup_request_id_reuse_not_polluted(build_adapter):
+    """After cleanup, reusing the same request_id must not be treated as finished."""
+    adapter, _ = build_adapter(stage_id=1)
+    req_id, ext_id = "req-reuse", "ext-reuse"
+    _populate_adapter_state(adapter, req_id, ext_id)
+
+    adapter.cleanup(req_id, ext_id)
+
+    assert req_id not in adapter.finished_requests
+    assert req_id not in adapter.get_req_chunk
+
+
+def test_cleanup_preserves_pending_save(build_adapter):
+    """Cleanup must NOT remove _pending_save_reqs to avoid losing unsent chunks."""
+    adapter, _ = build_adapter(stage_id=1)
+    req_id, ext_id = "req-4", "ext-4"
+    _populate_adapter_state(adapter, req_id, ext_id)
+
+    pending_task = {"put_key": f"{ext_id}_1_0", "data": {"x": 1}}
+    adapter._pending_save_reqs.append(pending_task)
+
+    adapter.cleanup(req_id, ext_id)
+
+    assert len(adapter._pending_save_reqs) == 1
+
+
+def test_cleanup_only_affects_target_request(build_adapter):
+    """Cleanup for one request must not affect another request's state."""
+    adapter, _ = build_adapter(stage_id=1)
+    _populate_adapter_state(adapter, "req-a", "ext-a")
+    _populate_adapter_state(adapter, "req-b", "ext-b")
+
+    adapter.cleanup("req-a", "ext-a")
+
+    assert "req-b" in adapter.finished_requests
+    assert "req-b" in adapter.get_req_chunk
+    assert "ext-b" in adapter.put_req_chunk
+    assert "ext-b" in adapter.request_payload
+    assert "ext-b" in adapter.code_prompt_token_ids
+    assert "req-b" in adapter.request_ids_mapping
+
+
+def test_cleanup_after_poll_flow(build_adapter):
+    """Simulate full load_async -> poll -> finished -> cleanup cycle."""
+    adapter, connector = build_adapter(stage_id=2, model_mode="ar")
+    request = _req("req-flow", RequestStatus.WAITING, external_req_id="ext-flow")
+
+    adapter.load_async(request)
+
+    adapter.request_ids_mapping["req-flow"] = "ext-flow"
+    payload = {"hidden_states": torch.tensor([[1.0]]), "finished": True}
+    connector.get.return_value = (payload, 8)
+    adapter._poll_single_request(request)
+
+    assert "req-flow" in adapter.finished_requests
+    assert adapter.get_req_chunk["req-flow"] == 1
+    assert "req-flow" in adapter.request_ids_mapping
+
+    adapter.cleanup("req-flow", "ext-flow")
+
+    assert "req-flow" not in adapter.finished_requests
+    assert "req-flow" not in adapter.get_req_chunk
+    assert "req-flow" not in adapter.request_ids_mapping
+    assert "ext-flow" not in adapter.request_payload
+
+
+# ---------------------------------------------------------------
+# Scheduler trigger tests
+# ---------------------------------------------------------------
+
+
+class _HashableRequest(SimpleNamespace):
+    """SimpleNamespace that can be added to a set (needed by scheduler internals)."""
+
+    def __hash__(self):
+        return hash(self.request_id)
+
+    def __eq__(self, other):
+        return getattr(other, "request_id", None) == self.request_id
+
+
+def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: MockerFixture):
+    """OmniGenerationScheduler must call adapter.cleanup when request finishes."""
+    cleanup_calls = []
+
+    adapter_mock = mocker.MagicMock()
+    adapter_mock.finished_requests = {"req-s1"}
+    adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
+
+    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+    scheduler = mocker.MagicMock()
+    scheduler.chunk_transfer_adapter = adapter_mock
+    scheduler.connector = None
+    scheduler.ec_connector = None
+    scheduler.perf_metrics = None
+    scheduler.log_stats = False
+    scheduler.recompute_kv_load_failures = False
+    scheduler.structured_output_manager = mocker.MagicMock()
+    scheduler.structured_output_manager.should_advance.return_value = False
+    scheduler.finished_req_ids_dict = {}
+    scheduler.kv_cache_manager.take_events.return_value = None
+    scheduler.kv_event_publisher = mocker.MagicMock()
+
+    request = _HashableRequest(
+        request_id="req-s1",
+        external_req_id="ext-s1",
+        status=RequestStatus.RUNNING,
+        is_finished=lambda: False,
+        num_computed_tokens=10,
+        num_prompt_tokens=10,
+        prompt_token_ids=list(range(10)),
+        num_output_placeholders=0,
+        sampling_params=None,
+        pooling_params=None,
+        stop_reason=None,
+        client_index=0,
+        take_events=lambda: [],
+        trace_headers=None,
+        num_cached_tokens=0,
+        num_external_computed_tokens=0,
+        num_nans_in_logits=0,
+        get_finished_reason=lambda: "stop",
+    )
+    scheduler.requests = {"req-s1": request}
+
+    scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
+    scheduler._free_request = mocker.MagicMock(return_value=None)
+    scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
+    scheduler.running = [request]
+    scheduler.waiting = mocker.MagicMock()
+    scheduler.waiting.remove_requests = mocker.MagicMock()
+    scheduler.make_stats = mocker.MagicMock(return_value=None)
+
+    scheduler_output = SimpleNamespace(
+        num_scheduled_tokens={"req-s1": 10},
+        scheduled_spec_decode_tokens={},
+        num_invalid_spec_tokens=0,
+    )
+    model_runner_output = SimpleNamespace(
+        sampled_token_ids=None,
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=None,
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        cudagraph_stats=None,
+        req_id_to_index={"req-s1": 0},
+    )
+
+    OmniGenerationScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
+
+    assert len(cleanup_calls) == 1
+    args, _ = cleanup_calls[0]
+    assert args[0] == "req-s1"
+    assert args[1] == "ext-s1"

tests/e2e/offline_inference/test_diffusion_cpu_offload.py

@@ -48,7 +48,6 @@ def inference(model_name: str, offload: bool = True):
 @pytest.mark.core_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"})
-@pytest.mark.skipif(current_omni_platform.is_npu() or current_omni_platform.is_rocm(), reason="Hardware not supported")
 @pytest.mark.parametrize("model_name", models)
 def test_cpu_offload_diffusion_model(model_name: str):
     try:

tests/e2e/offline_inference/test_diffusion_layerwise_offload.py

@@ -30,7 +30,7 @@ def run_inference(
 ) -> float:
     # For now, only support on GPU, so apply torch.cuda operations here
     # NPU / ROCm platforms are expected to be detected and skipped this test function
-    torch.cuda.empty_cache()
+    current_omni_platform.empty_cache()
     device_index = torch.cuda.current_device()
     monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
@@ -68,7 +68,6 @@ def run_inference(
     return peak
 
 
-@pytest.mark.skipif(current_omni_platform.is_npu() or current_omni_platform.is_rocm(), reason="Hardware not supported")
 @pytest.mark.parametrize("model_name", MODELS_SAVED_MEMORY_MB.keys())
 def test_layerwise_offload_diffusion_model(model_name: str):
     """Test that layerwise offloading reduces GPU memory usage.

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -71,7 +71,7 @@ def get_system_prompt():
 
 def get_prompt(prompt_type="text_only"):
     prompts = {
-        "text_only": "What is the capital of China?",
+        "text_only": "What is the capital of China? Answer in 20 words.",
         "mix": "What is recited in the audio? What is in this image? What is in this video?",
         "text_video": "What is in this video? ",
         "text_image": "What is in this image? ",

tests/e2e/stage_configs/qwen2_5_omni_ci.yaml

@@ -22,6 +22,7 @@ stage_args:
       trust_remote_code: true
       engine_output_type: latent
       enable_prefix_caching: false
+      mm_processor_cache_gb: 0
     is_comprehension: true
     final_output: true
     final_output_type: text

tests/e2e/stage_configs/qwen2_5_omni_thinker_ci.yaml

@@ -0,0 +1,31 @@
+stage_args:
+  - stage_id: 0
+    runtime:
+      process: true            # Run this stage in a separate process
+      devices: "0"            # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen2_5OmniForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      max_model_len: 16384
+      max_num_batched_tokens: 16384
+      max_num_seqs: 1
+      gpu_memory_utilization: 0.9
+      skip_mm_profiling: true
+      enforce_eager: true  # Now we only support eager mode
+      trust_remote_code: true
+      engine_output_type: latent
+      enable_prefix_caching: false
+    is_comprehension: true
+    final_output: true
+    final_output_type: text
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 128
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1

tests/e2e/stage_configs/qwen3_omni_thinker_ci.yaml

@@ -1,31 +0,0 @@
-# The following config has been verified on 2x H100-80G GPUs.
-stage_args:
-  - stage_id: 0
-    runtime:
-      devices: "0,1"
-      max_batch_size: 5
-    engine_args:
-      model_stage: thinker
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.6
-      enforce_eager: false
-      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
-      load_format: dummy
-    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: 100
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05

tests/engine/test_async_omni_engine_abort.py

@@ -15,8 +15,8 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 SEED = 42
 
-stage_config = str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen3_omni_thinker_ci.yaml")
-model = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+stage_config = str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen2_5_omni_thinker_ci.yaml")
+model = "Qwen/Qwen2.5-Omni-7B"
 
 
 async def generate(
@@ -60,7 +60,7 @@ async def generate(
 
 @pytest.mark.core_model
 @pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
 @pytest.mark.asyncio
 async def test_abort():
     with ExitStack() as after:

tests/entrypoints/openai_api/test_serving_speech.py

@@ -7,6 +7,7 @@ import pytest
 import torch
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
+from pydantic import ValidationError
 from pytest_mock import MockerFixture
 
 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
@@ -507,3 +508,110 @@ class TestTTSMethods:
         error = server._validate_tts_request(req)
         assert error is not None
         assert "max 10 characters" in error
+
+
+class TestStreamingProtocolValidation:
+    """Unit tests for the stream field validators in OpenAICreateSpeechRequest."""
+
+    def test_stream_validation_errors(self):
+        """stream=True requires response_format='pcm' and speed=1.0."""
+        with pytest.raises(ValidationError, match="response_format='pcm'"):
+            OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="wav")
+        with pytest.raises(ValidationError, match="Speed adjustment is not supported"):
+            OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm", speed=2.0)
+
+    def test_stream_valid(self):
+        """stream=True + response_format='pcm' + speed=1.0 is accepted."""
+        req = OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm")
+        assert req.stream is True
+
+    def test_sse_stream_format_is_blocked(self):
+        """stream_format='sse' is blocked."""
+        with pytest.raises(ValidationError, match="sse"):
+            OpenAICreateSpeechRequest(input="Hello", stream_format="sse")
+
+
+class TestStreamingResponse:
+    """Integration tests for the streaming audio response path."""
+
+    @pytest.fixture
+    def streaming_app(self, mocker: MockerFixture):
+        """Test app whose mock engine yields one intermediate chunk then a final chunk."""
+        mock_engine_client = mocker.MagicMock()
+        mock_engine_client.errored = False
+
+        def _make_output(finished: bool) -> OmniRequestOutput:
+            chunk = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 24000))
+
+            class MockCompletionOutput:
+                def __init__(self, index: int = 0):
+                    self.index = index
+                    self.text = ""
+                    self.token_ids = []
+                    self.finish_reason = "stop"
+                    self.stop_reason = None
+                    self.logprobs = None
+
+            class MockRequestOutput:
+                def __init__(self, audio_tensor: torch.Tensor):
+                    self.request_id = "speech-stream-test"
+                    self.outputs = [MockCompletionOutput(index=0)]
+                    self.multimodal_output = {"audio": audio_tensor}
+                    self.finished = finished
+                    self.prompt_token_ids = None
+                    self.encoder_prompt_token_ids = None
+                    self.num_cached_tokens = None
+                    self.prompt_logprobs = None
+                    self.kv_transfer_params = None
+
+            return OmniRequestOutput(
+                stage_id=0,
+                final_output_type="audio",
+                request_output=MockRequestOutput(audio_tensor=chunk),
+                finished=finished,
+            )
+
+        async def mock_generate_streaming(*args, **kwargs):
+            yield _make_output(finished=False)
+            yield _make_output(finished=True)
+
+        mock_engine_client.generate = mocker.MagicMock(side_effect=mock_generate_streaming)
+        mock_engine_client.default_sampling_params_list = [{}]
+        mock_models = mocker.MagicMock()
+        mock_models.is_base_model.return_value = True
+
+        speech_server = OmniOpenAIServingSpeech(
+            engine_client=mock_engine_client,
+            models=mock_models,
+            request_logger=mocker.MagicMock(),
+        )
+
+        original_create_speech = speech_server.create_speech
+        sig = signature(original_create_speech)
+        new_parameters = [p for name, p in sig.parameters.items() if name != "raw_request"]
+        new_sig = Signature(parameters=new_parameters, return_annotation=sig.return_annotation)
+
+        async def awaitable_create_speech(*args, **kwargs):
+            return await original_create_speech(*args, **kwargs)
+
+        awaitable_create_speech.__signature__ = new_sig
+        speech_server.create_speech = awaitable_create_speech
+
+        app = FastAPI()
+        app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None)
+        return app
+
+    def test_streaming(self, streaming_app):
+        """stream=True must return audio/pcm with non-empty body."""
+        client = TestClient(streaming_app)
+        response = client.post("/v1/audio/speech", json={"input": "Hello", "stream": True, "response_format": "pcm"})
+        assert response.status_code == 200
+        assert "audio/pcm" in response.headers["content-type"]
+        assert len(response.content) > 0
+
+    def test_non_streaming_unchanged(self, streaming_app):
+        """Non-streaming path must still return audio/wav."""
+        client = TestClient(streaming_app)
+        response = client.post("/v1/audio/speech", json={"input": "Hello", "response_format": "wav"})
+        assert response.status_code == 200
+        assert "audio/wav" in response.headers["content-type"]

tests/examples/online_serving/test_qwen2_5_omni.py

@@ -236,7 +236,6 @@ def test_modality_control_003(omni_server) -> None:
     # TODO: Verify the E2E latency after confirmation baseline.
 
 
-@pytest.mark.skip(reason="There is a known issue with stream error.")
 @pytest.mark.advanced_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})